# Copyright (c) 2026 Simulacra Research Inc. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import math from typing import NamedTuple, Sequence import equinox as eqx import jax import jax.numpy as jnp import numpy as np from hamiltonzero.inference import PreparedInference class BasisSamplerState(NamedTuple): bits: jax.Array log_abs: jax.Array key: jax.Array accepted: jax.Array proposed: jax.Array class Renyi2Result(NamedTuple): purity: float | None imaginary_mean: float | None standard_error: float | None renyi2_nats: float | None renyi2_bits: float | None resolved: bool failure_reasons: tuple[str, ...] naive_standard_error: float | None imaginary_standard_error: float | None imaginary_naive_standard_error: float | None integrated_autocorrelation_time_blocks: float | None integrated_autocorrelation_time_imaginary_blocks: float | None effective_blocks: float effective_imaginary_blocks: float largest_absolute_block_fraction: float | None renyi2_standard_error_nats: float | None renyi2_lower_3sigma_nats: float | None n_blocks: int mean_log_abs: float mean_phase: float block_log_abs: np.ndarray block_phase: np.ndarray swap_log_abs: np.ndarray swap_phase: np.ndarray valid_denominator: np.ndarray def _geometry(prepared: PreparedInference) -> tuple[jax.Array, int]: if not isinstance(prepared, PreparedInference): raise TypeError("prepared must be a PreparedInference") route_host = np.asarray(jax.device_get(prepared.route), dtype=np.int32) if route_host.ndim != 1: raise ValueError("prepared route must have shape [N]") if not np.array_equal(np.sort(route_host), np.arange(route_host.size)): raise ValueError("prepared route must be a permutation") mask = np.asarray(jax.device_get(prepared._initial_context.mask), dtype=np.bool_) n_spins = int(np.sum(mask)) if mask.shape != route_host.shape or not np.array_equal( mask, np.arange(mask.size) < n_spins ): raise ValueError("prepared physical sites must be a contiguous prefix") return jnp.asarray(prepared.route, dtype=jnp.int32), n_spins def _as_bits(bits, *, n_spins: int) -> jax.Array: value = jnp.asarray(bits) if ( value.ndim not in (2, 3) or any(size < 1 for size in value.shape[:-1]) or value.shape[-1] != n_spins ): raise ValueError( f"bits must have shape [pairs, {n_spins}] or [blocks, pairs, {n_spins}]" ) if value.dtype != jnp.bool_: invalid = jnp.any((value != 0) & (value != 1)) if bool(jax.device_get(invalid)): raise ValueError("bits must contain only zero and one") return value.astype(jnp.bool_) def _routed_corners(bits, route, *, n_spins: int): width = int(route.shape[0]) padding = width - int(n_spins) if padding < 0: raise ValueError("physical spin count exceeds compiled width") full_bits = jnp.pad(bits, ((0, 0), (0, padding)), constant_values=False) routed = jnp.take(full_bits, route, axis=-1) up = jnp.logical_not(routed).astype(jnp.float32) down = routed.astype(jnp.float32) zeros = jnp.zeros_like(up) return jnp.stack((up, zeros, down, zeros), axis=-1) def _basis_log_wavefunction(wavefunction, bits, route, *, n_spins: int): q = _routed_corners(bits, route, n_spins=n_spins) return wavefunction(q, None, 0.0) def _basis_log_abs(wavefunction, bits, route, *, n_spins: int): log_abs, _phase = _basis_log_wavefunction( wavefunction, bits, route, n_spins=n_spins ) return log_abs def _metropolis_log_acceptance(current_log_abs, proposed_log_abs): current_finite = jnp.isfinite(current_log_abs) proposed_finite = jnp.isfinite(proposed_log_abs) log_ratio = 2.0 * (proposed_log_abs - current_log_abs) log_accept = jnp.minimum(jnp.zeros_like(log_ratio), log_ratio) both_zero = jnp.isneginf(current_log_abs) & jnp.isneginf(proposed_log_abs) recover_to_finite = jnp.logical_not(current_finite) & proposed_finite invalid_proposal = jnp.logical_not(proposed_finite) & jnp.logical_not(both_zero) log_accept = jnp.where(both_zero | recover_to_finite, 0.0, log_accept) return jnp.where(invalid_proposal, -jnp.inf, log_accept) def _state_from_bits(key, wavefunction, bits, route, *, n_spins: int): log_abs = _basis_log_abs(wavefunction, bits, route, n_spins=n_spins) return BasisSamplerState( bits=bits, log_abs=log_abs, key=key, accepted=jnp.asarray(0, dtype=jnp.int32), proposed=jnp.asarray(0, dtype=jnp.int32), ) def _basis_step(state, wavefunction, route, *, n_spins: int): site_key, accept_key, next_key = jax.random.split(state.key, 3) batch_size = state.bits.shape[0] sites = jax.random.randint(site_key, (batch_size,), 0, n_spins, dtype=jnp.int32) rows = jnp.arange(batch_size) proposed_bits = state.bits.at[rows, sites].set( jnp.logical_not(state.bits[rows, sites]) ) proposed_log_abs = _basis_log_abs( wavefunction, proposed_bits, route, n_spins=n_spins ) log_accept = _metropolis_log_acceptance(state.log_abs, proposed_log_abs) log_uniform = jnp.log( jax.random.uniform(accept_key, state.log_abs.shape, dtype=state.log_abs.dtype) ) accept = log_uniform < log_accept return BasisSamplerState( bits=jnp.where(accept[:, None], proposed_bits, state.bits), log_abs=jnp.where(accept, proposed_log_abs, state.log_abs), key=next_key, accepted=state.accepted + jnp.sum(accept, dtype=jnp.int32), proposed=state.proposed + jnp.asarray(state.bits.shape[0], dtype=jnp.int32), ) @eqx.filter_jit def _run_basis_steps(state, wavefunction, route, *, n_spins: int, n_steps: int): def one_step(carry, _): return _basis_step(carry, wavefunction, route, n_spins=n_spins), None state, _ = jax.lax.scan(one_step, state, xs=None, length=n_steps) return state def _require_finite_state(state: BasisSamplerState) -> None: finite = np.asarray(jax.device_get(jnp.isfinite(state.log_abs))) if not np.all(finite): raise RuntimeError( "basis burn-in ended with a zero or non-finite wavefunction coefficient" ) def burn_in_basis( prepared: PreparedInference, key, *, batch_size: int = 256, burn_in: int = 1024, ): if int(batch_size) < 1: raise ValueError("batch_size must be positive") if int(burn_in) < 0: raise ValueError("burn_in must be non-negative") route, n_spins = _geometry(prepared) bits_key, state_key = jax.random.split(key) bits = jax.random.bernoulli(bits_key, shape=(int(batch_size), n_spins)) state = _state_from_bits( state_key, prepared.wavefunction, bits, route, n_spins=n_spins, ) state = _run_basis_steps( state, prepared.wavefunction, route, n_spins=n_spins, n_steps=int(burn_in), ) jax.block_until_ready(state.log_abs) _require_finite_state(state) return state, state.bits def step_basis( prepared: PreparedInference, state: BasisSamplerState, *, steps: int = 24, ): if not isinstance(state, BasisSamplerState): raise TypeError("state must be a BasisSamplerState") if int(steps) < 1: raise ValueError("steps must be positive") route, n_spins = _geometry(prepared) if state.bits.shape[-1] != n_spins: raise ValueError("basis state width does not match the prepared system") state = _run_basis_steps( state, prepared.wavefunction, route, n_spins=n_spins, n_steps=int(steps), ) jax.block_until_ready(state.log_abs) _require_finite_state(state) return state, state.bits def _subsystem_mask( subsystem: Sequence[int] | Sequence[bool] | np.ndarray, *, n_spins: int, ) -> jax.Array: value = np.asarray(subsystem) if value.ndim == 1 and value.size == 0: mask = np.zeros((n_spins,), dtype=np.bool_) elif value.dtype == np.bool_: if value.shape != (n_spins,): raise ValueError(f"boolean subsystem mask must have shape [{n_spins}]") mask = value else: if value.ndim != 1: raise ValueError("subsystem site indices must be one-dimensional") if not np.issubdtype(value.dtype, np.integer): raise TypeError("subsystem must contain integer sites or booleans") sites = value.astype(np.int64) if len(np.unique(sites)) != sites.size: raise ValueError("subsystem site indices must be unique") if np.any(sites < 0) or np.any(sites >= n_spins): raise ValueError("subsystem site index is out of range") mask = np.zeros((n_spins,), dtype=np.bool_) mask[sites] = True return jnp.asarray(mask) @eqx.filter_jit def _swap_log_ratios( wavefunction, replica_x, replica_y, route, region_mask, *, n_spins: int, ): swapped_x = jnp.where(region_mask[None, :], replica_y, replica_x) swapped_y = jnp.where(region_mask[None, :], replica_x, replica_y) denominator_x, phase_x = _basis_log_wavefunction( wavefunction, replica_x, route, n_spins=n_spins ) denominator_y, phase_y = _basis_log_wavefunction( wavefunction, replica_y, route, n_spins=n_spins ) numerator_x, numerator_phase_x = _basis_log_wavefunction( wavefunction, swapped_x, route, n_spins=n_spins ) numerator_y, numerator_phase_y = _basis_log_wavefunction( wavefunction, swapped_y, route, n_spins=n_spins ) valid_denominator = jnp.isfinite(denominator_x) & jnp.isfinite(denominator_y) log_abs = numerator_x + numerator_y - denominator_x - denominator_y phase = numerator_phase_x + numerator_phase_y - phase_x - phase_y phase = jnp.arctan2(jnp.sin(phase), jnp.cos(phase)) log_abs = jnp.where(valid_denominator, log_abs, -jnp.inf) phase = jnp.where(valid_denominator, phase, 0.0) exact_identity = jnp.logical_or( jnp.all(jnp.logical_not(region_mask)), jnp.all(region_mask) ) log_abs = jnp.where(exact_identity & valid_denominator, 0.0, log_abs) phase = jnp.where(exact_identity & valid_denominator, 0.0, phase) return log_abs, phase, valid_denominator def _complex_mean_log_polar(log_abs, phase) -> tuple[float, float]: logs = np.asarray(log_abs, dtype=np.float64).reshape(-1) phases = np.asarray(phase, dtype=np.float64).reshape(-1) if logs.shape != phases.shape or logs.size == 0: raise ValueError("log_abs and phase must have matching nonempty shapes") if ( np.any(np.isnan(logs)) or np.any(np.isposinf(logs)) or np.any(~np.isfinite(phases)) ): raise ValueError("non-finite SWAP log-polar sample") finite = np.isfinite(logs) if not np.any(finite): return -math.inf, 0.0 pivot = float(np.max(logs[finite])) scaled = np.zeros(logs.shape, dtype=np.complex128) scaled[finite] = np.exp(logs[finite] - pivot + 1j * phases[finite]) scaled_mean = np.mean(scaled) magnitude = float(abs(scaled_mean)) if magnitude == 0.0: return -math.inf, 0.0 return pivot + math.log(magnitude), float(np.angle(scaled_mean)) def _log_polar_to_complex(log_abs: float, phase: float) -> complex: if log_abs == -math.inf: return 0.0j if not math.isfinite(log_abs) or not math.isfinite(phase): raise ValueError("log-polar scalar must be finite or exact zero") if log_abs > math.log(np.finfo(np.float64).max): raise OverflowError("complex SWAP block mean exceeds float64 range") return complex(math.exp(log_abs) * np.exp(1j * phase)) def _integrated_autocorrelation_time(values) -> float: x = np.asarray(values, dtype=np.float64).reshape(-1) if x.size < 2: return 0.5 x = x - np.mean(x) variance = float(np.dot(x, x) / x.size) if not math.isfinite(variance) or variance <= 0.0: return 0.5 correlations = [] for lag in range(1, x.size): covariance = float(np.dot(x[:-lag], x[lag:]) / x.size) correlations.append(covariance / variance) tau = 0.5 previous_pair = math.inf for offset in range(0, len(correlations) - 1, 2): pair = correlations[offset] + correlations[offset + 1] if not math.isfinite(pair) or pair <= 0.0: break pair = min(pair, previous_pair) tau += pair previous_pair = pair return max(0.5, float(tau)) def _summarize_blocks(block_log_abs, block_phase): logs = np.asarray(block_log_abs, dtype=np.float64).reshape(-1) phases = np.asarray(block_phase, dtype=np.float64).reshape(-1) if logs.shape != phases.shape or logs.size == 0: raise ValueError("block log magnitudes and phases must match") try: blocks = np.asarray( [ _log_polar_to_complex(float(value), float(angle)) for value, angle in zip(logs, phases, strict=True) ], dtype=np.complex128, ) except (OverflowError, ValueError): return { "n_blocks": int(logs.size), "purity": None, "naive_standard_error": None, "standard_error": None, "imaginary_mean": None, "imaginary_naive_standard_error": None, "imaginary_standard_error": None, "integrated_autocorrelation_time_blocks": None, "integrated_autocorrelation_time_imaginary_blocks": None, "effective_blocks": 0.0, "effective_imaginary_blocks": 0.0, "largest_absolute_block_fraction": None, "resolved": False, "failure_reasons": ("block_mean_float64_overflow_or_nonfinite",), "renyi2_nats": None, "renyi2_bits": None, "renyi2_standard_error_nats": None, "renyi2_lower_3sigma_nats": None, } n_blocks = int(blocks.size) real = blocks.real imaginary = blocks.imag purity = float(np.mean(real)) imaginary_mean = float(np.mean(imaginary)) naive_standard_error = ( float(np.std(real, ddof=1) / math.sqrt(n_blocks)) if n_blocks > 1 else math.inf ) imaginary_naive_standard_error = ( float(np.std(imaginary, ddof=1) / math.sqrt(n_blocks)) if n_blocks > 1 else math.inf ) tau = _integrated_autocorrelation_time(real) effective_blocks = float(n_blocks / (2.0 * tau)) imaginary_tau = _integrated_autocorrelation_time(imaginary) effective_imaginary_blocks = float(n_blocks / (2.0 * imaginary_tau)) standard_error = ( float(np.std(real, ddof=1) / math.sqrt(effective_blocks)) if n_blocks > 1 else math.inf ) imaginary_standard_error = ( float(np.std(imaginary, ddof=1) / math.sqrt(effective_imaginary_blocks)) if n_blocks > 1 else math.inf ) absolute_sum = float(np.sum(np.abs(blocks))) tail_fraction = ( float(np.max(np.abs(blocks)) / absolute_sum) if absolute_sum > 0.0 else 1.0 ) failures = [] if n_blocks < 16: failures.append("too_few_blocks") if not math.isfinite(purity) or not math.isfinite(standard_error): failures.append("nonfinite_real_estimate") elif purity <= 3.0 * standard_error: failures.append("purity_not_resolved_above_zero") if math.isfinite(imaginary_standard_error): if abs(imaginary_mean) > 3.0 * imaginary_standard_error: failures.append("imaginary_null_test_failed") else: failures.append("nonfinite_imaginary_uncertainty") if purity > 1.0: failures.append("purity_point_above_physical_upper_bound") if effective_blocks < 8.0: failures.append("insufficient_effective_blocks") if tail_fraction > 0.25: failures.append("single_block_tail_dominance") resolved = not failures entropy = -math.log(purity) if resolved else None entropy_standard_error = standard_error / purity if resolved else None purity_upper_3sigma = min(1.0, purity + 3.0 * standard_error) if resolved else None entropy_lower_3sigma = ( max(0.0, -math.log(purity_upper_3sigma)) if purity_upper_3sigma is not None else None ) return { "n_blocks": n_blocks, "purity": purity, "naive_standard_error": naive_standard_error, "standard_error": standard_error, "imaginary_mean": imaginary_mean, "imaginary_naive_standard_error": imaginary_naive_standard_error, "imaginary_standard_error": imaginary_standard_error, "integrated_autocorrelation_time_blocks": tau, "integrated_autocorrelation_time_imaginary_blocks": imaginary_tau, "effective_blocks": effective_blocks, "effective_imaginary_blocks": effective_imaginary_blocks, "largest_absolute_block_fraction": tail_fraction, "resolved": resolved, "failure_reasons": tuple(failures), "renyi2_nats": entropy, "renyi2_bits": entropy / math.log(2.0) if entropy is not None else None, "renyi2_standard_error_nats": entropy_standard_error, "renyi2_lower_3sigma_nats": entropy_lower_3sigma, } def _evaluate_swap( prepared, x, y, route, mask, *, n_spins: int, chunk_size: int, ): logs = [] phases = [] valid = [] for start in range(0, x.shape[0], chunk_size): stop = min(start + chunk_size, x.shape[0]) values = _swap_log_ratios( prepared.wavefunction, x[start:stop], y[start:stop], route, mask, n_spins=n_spins, ) values = jax.device_get(values) logs.append(np.asarray(values[0])) phases.append(np.asarray(values[1])) valid.append(np.asarray(values[2])) return ( np.concatenate(logs), np.concatenate(phases), np.concatenate(valid), ) def _result(log_abs, phase, valid) -> Renyi2Result: swap_log_abs = np.asarray(log_abs) swap_phase = np.asarray(phase) valid_denominator = np.asarray(valid, dtype=np.bool_) if ( swap_log_abs.ndim != 2 or swap_log_abs.shape != swap_phase.shape or swap_log_abs.shape != valid_denominator.shape ): raise ValueError("SWAP blocks must have aligned shape [blocks, pairs]") if not np.all(valid_denominator): raise RuntimeError("SWAP denominator contains a zero wavefunction coefficient") block_values = [ _complex_mean_log_polar(logs, phases) for logs, phases in zip(swap_log_abs, swap_phase, strict=True) ] block_log_abs = np.asarray([value[0] for value in block_values], dtype=np.float64) block_phase = np.asarray([value[1] for value in block_values], dtype=np.float64) summary = _summarize_blocks(block_log_abs, block_phase) mean_log_abs, mean_phase = _complex_mean_log_polar(block_log_abs, block_phase) return Renyi2Result( purity=summary["purity"], imaginary_mean=summary["imaginary_mean"], standard_error=summary["standard_error"], renyi2_nats=summary["renyi2_nats"], renyi2_bits=summary["renyi2_bits"], resolved=summary["resolved"], failure_reasons=summary["failure_reasons"], naive_standard_error=summary["naive_standard_error"], imaginary_standard_error=summary["imaginary_standard_error"], imaginary_naive_standard_error=summary["imaginary_naive_standard_error"], integrated_autocorrelation_time_blocks=summary[ "integrated_autocorrelation_time_blocks" ], integrated_autocorrelation_time_imaginary_blocks=summary[ "integrated_autocorrelation_time_imaginary_blocks" ], effective_blocks=summary["effective_blocks"], effective_imaginary_blocks=summary["effective_imaginary_blocks"], largest_absolute_block_fraction=summary["largest_absolute_block_fraction"], renyi2_standard_error_nats=summary["renyi2_standard_error_nats"], renyi2_lower_3sigma_nats=summary["renyi2_lower_3sigma_nats"], n_blocks=summary["n_blocks"], mean_log_abs=mean_log_abs, mean_phase=mean_phase, block_log_abs=block_log_abs, block_phase=block_phase, swap_log_abs=swap_log_abs, swap_phase=swap_phase, valid_denominator=valid_denominator, ) def renyi2_purity( prepared: PreparedInference, replica_x, replica_y, subsystem: Sequence[int] | Sequence[bool] | np.ndarray, *, chunk_size: int = 256, ) -> Renyi2Result: if int(chunk_size) < 1: raise ValueError("chunk_size must be positive") route, n_spins = _geometry(prepared) x = _as_bits(replica_x, n_spins=n_spins) y = _as_bits(replica_y, n_spins=n_spins) if x.shape != y.shape: raise ValueError("replica batches must have the same shape") if x.ndim == 2: x = x[None, ...] y = y[None, ...] n_blocks, pairs_per_block = x.shape[:2] mask = _subsystem_mask(subsystem, n_spins=n_spins) values = _evaluate_swap( prepared, x.reshape((n_blocks * pairs_per_block, n_spins)), y.reshape((n_blocks * pairs_per_block, n_spins)), route, mask, n_spins=n_spins, chunk_size=int(chunk_size), ) return _result( values[0].reshape((n_blocks, pairs_per_block)), values[1].reshape((n_blocks, pairs_per_block)), values[2].reshape((n_blocks, pairs_per_block)), ) def measure_renyi2( prepared: PreparedInference, replica_x: BasisSamplerState, replica_y: BasisSamplerState, subsystem: Sequence[int] | Sequence[bool] | np.ndarray, *, blocks: int = 16, samples_per_block: int = 1, steps_between: int = 24, chunk_size: int = 256, ): if not isinstance(replica_x, BasisSamplerState) or not isinstance( replica_y, BasisSamplerState ): raise TypeError("replicas must be BasisSamplerState values") if int(blocks) < 1: raise ValueError("blocks must be positive") if int(samples_per_block) < 1: raise ValueError("samples_per_block must be positive") if int(steps_between) < 1: raise ValueError("steps_between must be positive") if int(chunk_size) < 1: raise ValueError("chunk_size must be positive") route, n_spins = _geometry(prepared) if replica_x.bits.shape != replica_y.bits.shape: raise ValueError("replica states must have the same walker shape") if replica_x.bits.ndim != 2 or replica_x.bits.shape[-1] != n_spins: raise ValueError("basis state width does not match the prepared system") _require_finite_state(replica_x) _require_finite_state(replica_y) mask = _subsystem_mask(subsystem, n_spins=n_spins) block_logs = [] block_phases = [] block_valid = [] for _ in range(int(blocks)): logs = [] phases = [] valid = [] for _ in range(int(samples_per_block)): replica_x = _run_basis_steps( replica_x, prepared.wavefunction, route, n_spins=n_spins, n_steps=int(steps_between), ) replica_y = _run_basis_steps( replica_y, prepared.wavefunction, route, n_spins=n_spins, n_steps=int(steps_between), ) values = _evaluate_swap( prepared, replica_x.bits, replica_y.bits, route, mask, n_spins=n_spins, chunk_size=int(chunk_size), ) logs.append(values[0]) phases.append(values[1]) valid.append(values[2]) block_logs.append(np.concatenate(logs)) block_phases.append(np.concatenate(phases)) block_valid.append(np.concatenate(valid)) result = _result( np.stack(block_logs), np.stack(block_phases), np.stack(block_valid), ) return replica_x, replica_y, result __all__ = [ "BasisSamplerState", "Renyi2Result", "burn_in_basis", "measure_renyi2", "renyi2_purity", "step_basis", ]