from pathlib import Path
import sys

import equinox as eqx
import jax
import jax.numpy as jnp
import networkx as nx
import numpy as np

from torx.psc import HybridPCircuit

ROOT = Path.cwd()
if not (ROOT / "helpers").exists() and (ROOT.parent / "helpers").exists():
    ROOT = ROOT.parent
HELPER_DIR = ROOT / "helpers"
sys.path.insert(0, str(HELPER_DIR))

from _notebook_paths import figure_dir
from _notebook_style import (
    apply_notebook_style,
    display_and_close,
    make_savefig,
)
from _langevin import (
    LangevinGate,
    boltzmann_1d_reference,
    ks_distance,
    langevin_theta,
    plot_ks_vs_step,
    plot_mala_reference,
    plot_site_mean_parity,
    plot_ula_bias,
)
import _plots_sampling as P_samp
import _plots_schematics as P_sch

FIGURE_DIR = figure_dir(ROOT)
apply_notebook_style()

SEED = 31

savefig = make_savefig(FIGURE_DIR)

graph = nx.cycle_graph(10)
# Add three long-range chords so the cycle has nonlocal couplings.
graph.add_edges_from([(0, 5), (2, 7), (3, 8)])
pos = nx.spring_layout(graph, seed=SEED)

n = graph.number_of_nodes()
adj = nx.to_numpy_array(graph, dtype=float)
# Antisymmetric field: a per-site tilt from negative to positive that sums to ~0.
field = np.linspace(-0.30, 0.30, n)

assert graph.number_of_nodes() == 10
assert graph.number_of_edges() == 13  # 10 cycle edges + 3 cross-chords
print(f"graph: {n} nodes, {graph.number_of_edges()} edges")

graph: 10 nodes, 13 edges

class LangevinGate(AbstractContinuousGate[dict[str, Array], tuple[int, ...]]):
    ...
    metropolis: bool = eqx.field(static=True, default=False)

    def energy(
        self, x: Float[Array, " d"], theta: dict[str, Array]
    ) -> Float[Array, ""]:
        """Soft-spin quartic Ising energy $V(x)$ that the gate descends."""
        s = jnp.tanh(x)
        coupling = -0.5 * theta["beta"] * (s @ (theta["adj"] @ s))
        bias = -(theta["field"] @ s)
        quartic = 0.25 * theta["lam_quartic"] * jnp.sum((x * x - 1.0) ** 2)
        return coupling + bias + quartic

    def sample(
        self,
        key: Key[Array, ""],
        inputs: Mapping[str, PyTree[Array]],
        params: dict[str, Array],
        ...
    ):
        """Advance the continuous state by one (optionally adjusted) Langevin step."""
        x = inputs["continuous"]
        eps = params["step"]
        temperature = params["temperature"]
        grad_x = jax.grad(self.energy)(x, params)

        noise_key, accept_key = jax.random.split(key)
        ...
        noise = jax.random.normal(noise_key, x.shape, dtype=x.dtype)
        ...

        proposal = x - eps * grad_x + jnp.sqrt(2.0 * temperature * eps) * noise

        if not self.metropolis:
            output = proposal
        else:
            grad_prop = jax.grad(self.energy)(proposal, params)
            # log target ratio: -(V(x') - V(x)) / T
            log_target = (
                -(self.energy(proposal, params) - self.energy(x, params)) / temperature
            )
            # log proposal ratio log q(x | x') - log q(x' | x); the Gaussian
            # proposal has mean y - eps * grad V(y) and covariance 2 T eps I.
            fwd = proposal - (x - eps * grad_x)
            bwd = x - (proposal - eps * grad_prop)
            log_proposal = (jnp.sum(fwd**2) - jnp.sum(bwd**2)) / (
                4.0 * temperature * eps
            )
            accept_prob = jnp.exp(jnp.minimum(0.0, log_target + log_proposal))
            ...
            accepted = jax.random.bernoulli(accept_key, accept_prob)
            ...
            output = jnp.where(accepted, proposal, x)

BETA = 0.35
LAM_QUARTIC = 0.6
STEP_SIZE = 0.04
TEMPERATURE = 0.35

DEPLOY_REPS = 1200     # Metropolis-adjusted steps per deployed chain
LONGRUN_REPS = 6000    # a much longer MALA run, used as a convergence reference
NUM_SAMPLES = 1024

theta = langevin_theta(
    adj=adj,
    field=field,
    beta=BETA,
    lam_quartic=LAM_QUARTIC,
    step=STEP_SIZE,
    temperature=TEMPERATURE,
)

def terminal_samples(theta, *, metropolis, reps, num_samples, dim, key):
    """Run scanned Torx transitions for independent Langevin chains."""
    gate = LangevinGate(sites=0, dims=(dim,), metropolis=metropolis)
    circuit = HybridPCircuit([gate])
    initial_state = {
        "discrete": jnp.zeros((0,), dtype=jnp.int32),
        "continuous": jnp.zeros(dim),
    }

    def run_chain(chain_key):
        def step(carry, _):
            state, step_key = carry
            step_key, sample_key = jax.random.split(step_key)
            state = circuit.sample(sample_key, state, [theta])
            return (state, step_key), None

        (state, _), _ = jax.lax.scan(
            step, (initial_state, chain_key), None, length=reps
        )
        return state

    keys = jax.random.split(key, num_samples)
    sampled = eqx.filter_jit(jax.vmap(run_chain))(keys)
    return np.asarray(sampled["continuous"])

fig = P_sch.draw_pcircuit(
    [("MALA", [0])],
    wire_labels=[rf"$x\in\mathbb{{R}}^{{{n}}}$"],
    title="One MALA step on the full spin vector",
    reps=DEPLOY_REPS,
)
display_and_close(fig)

LAM_1D = 1.5      # a stiffer well than the graph run, to make ULA's bias vivid
T_1D = 0.30
H_1D = 0.25
T_SIM = 60.0      # fixed simulated time; reps = T_sim / step isolates step-size bias
N_1D = 8000

EPS_SWEEP = [0.02, 0.06, 0.11, 0.16]
EPS_MATCH = 0.06  # step for the MALA-matches-reference panel
EPS_BIAS = 0.16   # step for the ULA-bias panel (the coarsest in the sweep)

grid = np.linspace(-3.0, 3.0, 1401)

ref_gate = LangevinGate(sites=0, dims=(1,), metropolis=True)
theta_1d = langevin_theta(
    adj=[[0.0]],
    field=[H_1D],
    beta=0.0,
    lam_quartic=LAM_1D,
    step=EPS_MATCH,
    temperature=T_1D,
)
ref_density, ref_cdf = boltzmann_1d_reference(ref_gate, theta_1d, grid)

ula_1d, mala_1d = {}, {}
ks_ula, ks_mala = [], []
for eps in EPS_SWEEP:
    reps = round(T_SIM / eps)  # fixed simulated time T_sim = reps * eps
    theta_eps = langevin_theta(
        adj=[[0.0]],
        field=[H_1D],
        beta=0.0,
        lam_quartic=LAM_1D,
        step=eps,
        temperature=T_1D,
    )
    ula_1d[eps] = terminal_samples(
        theta_eps, metropolis=False, reps=reps, num_samples=N_1D, dim=1,
        key=jax.random.key(SEED + 100),
    )
    mala_1d[eps] = terminal_samples(
        theta_eps, metropolis=True, reps=reps, num_samples=N_1D, dim=1,
        key=jax.random.key(SEED + 200),
    )
    ks_ula.append(ks_distance(ula_1d[eps], grid, ref_cdf))
    ks_mala.append(ks_distance(mala_1d[eps], grid, ref_cdf))

for eps, ku, km in zip(EPS_SWEEP, ks_ula, ks_mala):
    print(f"eps={eps:.2f}  reps={round(T_SIM / eps):5d}  ULA KS={ku:.4f}   MALA KS={km:.4f}")

ks_match = ks_mala[EPS_SWEEP.index(EPS_MATCH)]
ks_bias_ula = ks_ula[EPS_SWEEP.index(EPS_BIAS)]
ks_bias_mala = ks_mala[EPS_SWEEP.index(EPS_BIAS)]

assert ks_match < 0.03, f"MALA does not match the reference: KS {ks_match:.4f}"
assert ks_bias_ula > 0.05, f"ULA bias not visible at eps={EPS_BIAS}: KS {ks_bias_ula:.4f}"
assert ks_bias_ula > 4 * ks_bias_mala, "MALA did not remove ULA's step-size bias"
assert ks_ula[-1] > ks_ula[0], "ULA KS should grow with the step size"

eps=0.02  reps= 3000  ULA KS=0.0100   MALA KS=0.0083
eps=0.06  reps= 1000  ULA KS=0.0199   MALA KS=0.0068
eps=0.11  reps=  545  ULA KS=0.0369   MALA KS=0.0080
eps=0.16  reps=  375  ULA KS=0.0805   MALA KS=0.0099

fig = plot_mala_reference(
    grid, ref_density, mala_1d[EPS_MATCH], ks=ks_match, temperature=T_1D
)
savefig(fig, "12_langevin_mala_reference")

fig = plot_ula_bias(
    grid,
    ref_density,
    ula_1d[EPS_BIAS],
    mala_1d[EPS_BIAS],
    step=EPS_BIAS,
    ks_ula=ks_bias_ula,
    ks_mala=ks_bias_mala,
)
savefig(fig, "12_langevin_ula_bias")

fig = plot_ks_vs_step(EPS_SWEEP, ks_ula, ks_mala)
savefig(fig, "12_langevin_ks_vs_step")

samples = terminal_samples(
    theta, metropolis=True, reps=DEPLOY_REPS, num_samples=NUM_SAMPLES, dim=n,
    key=jax.random.key(SEED),
)
assert samples.shape == (NUM_SAMPLES, n), f"expected ({NUM_SAMPLES}, {n}), got {samples.shape}"
assert np.all(np.isfinite(samples)), "non-finite samples"
print(f"samples shape: {samples.shape}  min={samples.min():.3f}  max={samples.max():.3f}")

samples shape: (1024, 10)  min=-2.173  max=2.175

long_samples = terminal_samples(
    theta, metropolis=True, reps=LONGRUN_REPS, num_samples=NUM_SAMPLES, dim=n,
    key=jax.random.key(SEED + 1),
)
mu_deploy = np.mean(np.tanh(samples), axis=0)
mu_long = np.mean(np.tanh(long_samples), axis=0)

# Monte Carlo standard error for each site's two independent mean estimates.
se_deploy = np.std(np.tanh(samples), axis=0, ddof=1) / np.sqrt(NUM_SAMPLES)
se_long = np.std(np.tanh(long_samples), axis=0, ddof=1) / np.sqrt(NUM_SAMPLES)
se_difference = np.sqrt(se_deploy**2 + se_long**2)
site_differences = np.abs(mu_deploy - mu_long)
standardized_residuals = site_differences / se_difference
worst_site = int(np.argmax(standardized_residuals))
conv_err = float(site_differences.max())
max_standardized_residual = float(standardized_residuals[worst_site])
print(
    f"largest |Δ_i|/SE_i = {max_standardized_residual:.2f} at site {worst_site}, "
    f"|Δ_i|={site_differences[worst_site]:.4f}, SE_i={se_difference[worst_site]:.4f}"
)
assert conv_err < 0.06, f"deployed MALA differs from long run: max abs diff {conv_err:.4f}"

largest |Δ_i|/SE_i = 1.16 at site 6, |Δ_i|=0.0347, SE_i=0.0301

fig = plot_site_mean_parity(
    mu_deploy,
    mu_long,
    se_short=se_deploy,
    se_long=se_long,
    reps_short=DEPLOY_REPS,
    reps_long=LONGRUN_REPS,
)
savefig(fig, "12_langevin_convergence_parity")

mean_soft_spin = np.mean(np.tanh(samples), axis=0)
terminal_magnetization = np.mean(np.tanh(samples), axis=1)
mean_mag = float(terminal_magnetization.mean())

fig = P_sch.plot_terminal_field(graph, pos, mean_soft_spin)
savefig(fig, "12_langevin_terminal_field")

fig = P_samp.plot_magnetization_histogram_langevin(
    terminal_magnetization,
    title_hist=f"Magnetization distribution  ({NUM_SAMPLES} paths)",
)
savefig(fig, "12_langevin_magnetization")

def run_mala_path(theta, *, dim, steps, key):
    """Record one directly sampled MALA path and its acceptance rate."""
    gate = LangevinGate(sites=0, dims=(dim,), metropolis=True)

    def body(carry, _):
        x, k = carry
        k, subkey = jax.random.split(k)
        substate = {"discrete": jnp.zeros((0,), dtype=jnp.int32), "continuous": x}
        x_next = gate.sample(subkey, substate, theta)
        moved = jnp.any(jnp.abs(x_next - x) > 0.0)  # a rejection leaves x unchanged
        return (x_next, k), (x, moved)

    # One compiled scan records intermediate states instead of only the terminal state.
    _, (path, moved) = eqx.filter_jit(
        lambda key: jax.lax.scan(body, (jnp.zeros(dim), key), None, length=steps)
    )(key)
    energies = jax.vmap(lambda x: gate.energy(x, theta))(path)
    return np.asarray(energies), float(np.mean(np.asarray(moved)))

TRACE_STEPS = 500
energies_arr, accept_rate = run_mala_path(
    theta, dim=n, steps=TRACE_STEPS, key=jax.random.key(SEED + 2)
)
assert np.all(np.isfinite(energies_arr)), "non-finite energy trace"

mean_energy = float(energies_arr.mean())
initial_energy = float(energies_arr[0])
tail_energy = float(energies_arr[-100:].mean())  # late-trace level shown in the figure and summary

fig = P_samp.plot_energy_trace(energies_arr, tail_energy)
savefig(fig, "12_langevin_energy_trace")

assert accept_rate > 0.5, f"MALA acceptance too low: {accept_rate:.3f}"
assert np.isfinite(mean_energy), "mean energy is non-finite"
assert energies_arr.min() < initial_energy, "energy never descended below initial value"
assert tail_energy < initial_energy - 0.1 * abs(initial_energy), (
    f"energy did not relax below the start: last-window mean {tail_energy:.3f} vs "
    f"initial {initial_energy:.3f}"
)

print(f"-- Langevin graph Ising  (n={n}, deploy reps={DEPLOY_REPS}, {NUM_SAMPLES} paths) --")
print(f"  MALA acceptance rate  : {accept_rate:.3f}")
print(
    f"  deployed vs long-run  : max |Δ_i|/SE_i = {max_standardized_residual:.2f} "
    f"(site {worst_site})"
)
print(f"  mean magnetization    : {mean_mag:+.4f}   (empirical, range [-1, 1])")
print(
    f"  magnetization range   : [{terminal_magnetization.min():+.3f}, {terminal_magnetization.max():+.3f}]"
)
print(f"  initial energy V(x_0) : {initial_energy:+.4f}")
print(f"  last-100-step mean    : {tail_energy:+.4f}  (relaxation diagnostic)")
print(f"  1-D MALA KS (eps={EPS_MATCH}): {ks_match:.4f}")
print(f"  1-D ULA  KS (eps={EPS_BIAS}): {ks_bias_ula:.4f}   (step-size bias MALA removes)")
print("--------------------------------------------------------")

-- Langevin graph Ising  (n=10, deploy reps=1200, 1024 paths) --
  MALA acceptance rate  : 0.974
  deployed vs long-run  : max |Δ_i|/SE_i = 1.16 (site 6)
  mean magnetization    : +0.0029   (empirical, range [-1, 1])
  magnetization range   : [-0.866, +0.844]
  initial energy V(x_0) : +1.5000
  last-100-step mean    : +0.1910  (relaxation diagnostic)
  1-D MALA KS (eps=0.06): 0.0068
  1-D ULA  KS (eps=0.16): 0.0805   (step-size bias MALA removes)
--------------------------------------------------------

Sampling Langevin dynamics on a graph Ising model¶

Setup¶

The graph and soft-spin Ising energy¶

The Langevin gate¶

Validating against an exact reference¶

Sampling the full graph¶

Terminal field and magnetization¶

Energy relaxation diagnostic¶

Verification¶

Conclusion¶

References¶