Sampling benchmarks and performance

Sampling benchmarks and performance#

What’s in this notebook? This advanced notebook collects sampler comparisons, bounded-search scaling, solver comparisons and production-style tuning. Use it when deciding which method to use and how to tune it after the basic ISD workflow is already familiar.

Runtime note: benchmark cells are disabled by default through RUN_BENCHMARKS = False. With the flag off, the notebook documents the comparison design and executes only setup/skip cells; turn the flag on only when you are prepared for substantial JAX compilation time and memory use.

Outline#

  1. Setup

  2. Method comparison and benchmarks

  3. Advanced: performance tuning

  4. Take-aways

  5. Further reading

Setup#

import warnings, time
from pathlib import Path
import numpy as np
from scipy.optimize import root

import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)

import matplotlib.pyplot as plt
import seaborn as sn
cmap = sn.color_palette("viridis", as_cmap=True)

import jaxvacua as jvc
from stringjax_tools import vmapping_func

warnings.filterwarnings("ignore")

RUN_BENCHMARKS = False

h12 = 2
model = jvc.FluxVacuaFinder(h12=h12, Q=276, model_ID=1, maximum_degree=2, map_to_fd=True)
sampler = jvc.data_sampler(model, flux_bounds=[-5, 5], moduli_bounds=[0, 10], axion_bounds=[-0.5, 0.5])
DW_batch = model.DW  # JAXVacua auto-vectorises paired leading-axis batches.
rns_key = jax.random.PRNGKey(42)
n_fl = model.n_fluxes
Nmax = 200

DATA_CANDIDATES = [
    Path("."),
    Path("../02_vacuum_finding"),
    Path("02_vacuum_finding"),
    Path("documentation/source/notebooks/02_vacuum_finding"),
]

def first_existing(filename):
    for base in DATA_CANDIDATES:
        candidate = base / filename
        if candidate.exists():
            return candidate
    return None

DATASET_B_PATH = first_existing("dataset_B.p")

print("RUN_BENCHMARKS =", RUN_BENCHMARKS)
print("Set RUN_BENCHMARKS = True to execute the expensive comparison cells.")
print("DATASET_B_PATH =", DATASET_B_PATH)

Resource guidance#

This notebook is intentionally not a quick-start tutorial. The benchmark cells can compile several JAX kernels, allocate large candidate batches, and run for minutes or longer on a laptop. Reduce n_manual, n_batch, max_batches, n_sample, and the number of sweep points first when experimenting locally. The default path keeps the methodology visible while avoiding accidental long runs.

Method comparison and benchmarks#

This section compares the methods exposed in NB07 and NB08 on the same canonical fixture. The goal is not to produce universal timing constants, but to show which trade-offs matter: completeness versus stochastic yield, Newton versus scipy refinement, and how aggressively one should sample moduli and flux candidates.

Bounding-box scaling with \(N_{\max}\)#

The eigenvalue-based bounding box from arXiv:2501.03984 defines the search region for NS-NS fluxes. The Cartesian product over the full \(h\)-vector grows as \(\sim N_{\max}^3\), making systematic enumeration infeasible beyond \(N_{\max} \sim 30\).

if RUN_BENCHMARKS:
    # Show how bounding box scales with Nmax
    dil_bounds = (np.sqrt(3)/2, 10.)
    moduli_bounds = (2., 5.)
    axion_bounds = (-0.5, 0.5)

    Nmax_values = [5, 10, 20, 50, 100, 200, 500]
    box_sizes = []

    for Nm in Nmax_values:
        s = jvc.data_sampler(model, moduli_bounds=moduli_bounds,
            dilaton_bounds=dil_bounds, axion_bounds=axion_bounds, seed=42)
        bf = jvc.flux_bounding.bounded_fluxes(model, sampler=s, Nmax=Nm)
        bf.compute_eigenvalue_bounds(200, verbose=False)
        dim = bf.dimension_H3
        h1m = int(np.ceil(bf._h1_box))
        h2m = int(np.ceil(bf._h2_box))
        n_box = (2*h1m+1)**dim * (2*h2m+1)**dim
        box_sizes.append(n_box)
        feasible = "✓ enumerate" if n_box < 1e8 else ("⚠ streaming" if n_box < 1e10 else "✗ sample only")
        print(f"Nmax={Nm:>3}: h1_max={h1m:>3}, h2_max={h2m:>2}, "
              f"box={n_box:>18,}  {feasible}")

    fig, ax = plt.subplots(figsize=(8, 4),dpi=150)
    ax.semilogy(Nmax_values, box_sizes, 'o-', color='steelblue', linewidth=2)
    ax.axhline(1e8, color='green', linestyle='--', alpha=0.5, label='enumerate feasible (<10⁸)')
    ax.axhline(1e10, color='orange', linestyle='--', alpha=0.5, label='streaming feasible (<10¹⁰)')
    ax.set_xlabel(r'$N_{\max}$')
    ax.set_ylabel('Bounding box size')
    ax.set_title('Search space growth with tadpole bound')
    ax.legend()
    plt.tight_layout()
    plt.show()
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

The four sampling methods#

Method

Approach

Completeness

Best regime

Manual ISD

Random \((z, \tau, h)\) → ISD complete \(f\) → Newton

None

Baseline only

enumerate_fluxes

Systematic: all \(h\) in bounding box, stream + filter

Complete

\(N_{\max} \lesssim 30\)

sample_bounded_fluxes (Gaussian)

\(h \sim \mathcal{N}(0, \sigma^2 M)\) → ISD → Newton

Stochastic

\(N_{\max} \gtrsim 50\)

sample_bounded_fluxes (Gaussian + LS)

Same prior + linearised_shifts (moves moduli)

Stochastic

Medium \(N_{\max}\)

Why the Gaussian prior \(h \sim \mathcal{N}(0, \sigma^2 M)\)?#

The continuous ISD tadpole is \(N_{\rm flux}^{\rm cont} = s \cdot h^T M^{-1} h\). For \(h\) drawn from \(\mathcal{N}(0, \sigma^2 M)\):

(1)#\[\mathbb{E}\bigl[s_{\min} \cdot h^T M^{-1} h\bigr] = s_{\min} \cdot \sigma^2 \cdot \mathrm{tr}(I) = s_{\min} \cdot \sigma^2 \cdot d\]

Choosing \(\sigma^2 = N_{\max} / (d \cdot s_{\min})\) ensures most samples satisfy \(N_{\rm flux} \lesssim N_{\max}\). The \(M\)-weighting concentrates samples along eigendirections where \(M^{-1}\) has small eigenvalues (low tadpole cost), matching the physical vacuum distribution.

Benchmark helper#

The helper below is slow because it deliberately compares several qualitatively different algorithms under one interface. It samples candidate moduli and fluxes, calls ISD completion, may enumerate bounded fluxes, and then Newton-refines candidates. Keep the benchmark flags off during a normal documentation run.

def run_benchmark(Nmax, moduli_bounds=(2., 5.), dil_bounds=(np.sqrt(3)/2, 10.),
                  axion_bounds=(-0.5, 0.5), n_manual=5000, max_batches_sample=5,
                  n_batch_sample=10_000, skip_enumerate=False):
    """Run all four methods and return a results dict."""
    
    sampler = jvc.data_sampler(model, moduli_bounds=moduli_bounds,
        dilaton_bounds=dil_bounds, axion_bounds=axion_bounds, seed=42)
    results = {}
    
    # ── Method 1: Manual ISD ──
    rng = np.random.default_rng(42)
    t0 = time.perf_counter()
    mod_pts = jnp.array(sampler.get_complex_moduli(n_manual), dtype=complex)
    tau_pts = jnp.array(sampler.get_complex_tau(n_manual), dtype=complex)
    h_all = rng.integers(-5, 6, (n_manual, model.n_fluxes))
    
    good = []
    for i in range(n_manual):
        h_i = jnp.array(h_all[i], dtype=float)
        if jnp.all(h_i == 0): continue
        try:
            fl = sampler.ISD_sampling(mod_pts[i], jnp.conj(mod_pts[i]),
                                       tau_pts[i], jnp.conj(tau_pts[i]),
                                       h_i, mode="H")
            if fl is not None and jnp.all(jnp.isfinite(fl)):
                fl_r = jnp.array(fl).real
                fl_r = fl_r.at[:n_fl].set(jnp.round(fl_r[:n_fl]))
                tad = abs(float(jnp.real(model.tadpole(fl_r))))
                if 0 < tad <= Nmax:
                    good.append((mod_pts[i], tau_pts[i], fl_r))
        except: pass
    
    n_conv = 0
    if good:
        bf_tmp = jvc.flux_bounding.bounded_fluxes(model, sampler=sampler, Nmax=Nmax)
        m_arr = jnp.array([x[0] for x in good], dtype=complex)
        t_arr = jnp.array([x[1] for x in good], dtype=complex)
        f_arr = jnp.array([x[2] for x in good])
        m_o, t_o, r_o = bf_tmp.newton_refine_batch(m_arr, t_arr, f_arr,
            step_size=1.0, tol=1e-10, max_iters=100)
        r_np = np.asarray(r_o)
        in_p = np.asarray(bf_tmp.in_patch_batch(m_o, t_o))
        n_conv = int(np.sum((r_np < 1e-10) & in_p))
    dt = time.perf_counter() - t0
    results['Manual ISD'] = {'vacua': n_conv, 'time': dt,
                              'isd_yield': f"{len(good)}/{n_manual}"}
    
    # ── Method 2: enumerate_fluxes ──
    if not skip_enumerate:
        bf_e = jvc.flux_bounding.bounded_fluxes(model, sampler=sampler, Nmax=Nmax)
        t0 = time.perf_counter()
        res_e = bf_e.enumerate_fluxes(
            n_sample=200, n_isd_per_h=10, max_h_candidates=5_000_000,
            refine=True, return_moduli=True, newton_tol=1e-10,
            confirm_streaming=False, verbose=False, n_moduli_batches=1)
        dt = time.perf_counter() - t0
        results['enumerate_fluxes'] = {'vacua': len(res_e), 'time': dt}
    
    # ── Method 3: sample_bounded (Gaussian, no LS) ──
    bf_s = jvc.flux_bounding.bounded_fluxes(model, sampler=sampler, Nmax=Nmax)
    t0 = time.perf_counter()
    res_s = bf_s.sample_bounded_fluxes(
        n_target=500, n_batch=n_batch_sample, n_sample=100, n_mod=10,
        max_batches=max_batches_sample, refine=True, return_moduli=True,
        newton_tol=1e-10, verbose=False)
    dt = time.perf_counter() - t0
    results['sample (Gaussian)'] = {'vacua': len(res_s), 'time': dt}
    
    # ── Method 4: sample_bounded (Gaussian + linearised_shifts) ──
    bf_ls = jvc.flux_bounding.bounded_fluxes(model, sampler=sampler, Nmax=Nmax)
    t0 = time.perf_counter()
    res_ls = bf_ls.sample_bounded_fluxes(
        n_target=500, n_batch=min(n_batch_sample, 5000), n_sample=100, n_mod=10,
        max_batches=max_batches_sample, refine=True, return_moduli=True,
        newton_tol=1e-10, verbose=False,
        use_linearised_shifts=True, n_isd_iters=5, n_moduli_batches=1)
    dt = time.perf_counter() - t0
    results['sample (Gauss+LS)'] = {'vacua': len(res_ls), 'time': dt}
    
    return results


def print_results(results, Nmax):
    """Pretty-print benchmark results."""
    print(f"\n{'='*65}")
    print(f"  Nmax = {Nmax}")
    print(f"{'='*65}")
    print(f"{'Method':<28} {'Vacua':>7} {'Time':>9} {'Rate':>10}")
    print(f"{'-'*65}")
    for name, r in results.items():
        n = r['vacua']; dt = r['time']
        rate = n / max(dt, 0.1)
        if dt < 120: ts = f"{dt:.1f}s"
        elif dt < 3600: ts = f"{dt/60:.1f}m"
        else: ts = f"{dt/3600:.1f}h"
        extra = f"  (ISD: {r['isd_yield']})" if 'isd_yield' in r else ""
        print(f"{name:<28} {n:>7} {ts:>9} {rate:>9.2f}/s{extra}")
    print(f"{'-'*65}")

Benchmark: \(N_{\max} = 10\) (small tadpole)#

At small \(N_{\max}\), the bounding box is manageable and enumerate_fluxes can systematically check all candidates. The tadpole pre-filter (continuous quadratic form \(s_{\min} h^T M^{-1} h \leq N_{\max}\)) eliminates most of the box, making enumeration practical.

When this benchmark is enabled, reported vacua are Newton-refined to \(\sum_I |D_I W| < 10^{-10}\) and verified to lie inside the sampler’s moduli patch.

if RUN_BENCHMARKS:
    res_10 = run_benchmark(Nmax=5, n_manual=5000, max_batches_sample=5, n_batch_sample=10_000)
    print_results(res_10, 10)
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Benchmark: \(N_{\max} = 200\) (large tadpole)#

At large \(N_{\max}\), the bounding box is too large for complete enumeration in this model. This is the regime where stochastic bounded sampling becomes the relevant tool, but the runtime depends strongly on n_batch, max_batches, n_sample, and the chosen refinement policy.

if RUN_BENCHMARKS:
    res_200 = run_benchmark(Nmax=200, skip_enumerate=True,
                            n_manual=5000, max_batches_sample=5, n_batch_sample=10_000)
    print_results(res_200, 200)
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Gaussian prior analysis#

The Gaussian prior \(h \sim \mathcal{N}(0, \sigma^2 M)\) is designed to concentrate samples where the continuous ISD tadpole is small. The disabled cell below compares that prior to uniform box sampling and, when the local Dataset B file is available, overlays the stored \(h\)-norm distribution as a sanity check rather than as a proof of exact recovery.

if RUN_BENCHMARKS:
    # Compare ||h|| distributions: Gaussian prior vs Dataset B
    Nmax_plot = 10
    s_min_plot = np.sqrt(3) / 2
    sampler_plot = jvc.data_sampler(model, moduli_bounds=(2.,5.),
        dilaton_bounds=(s_min_plot, 10.), axion_bounds=(-0.5, 0.5), seed=42)

    # Representative ISD matrix
    z_rep = jnp.array(sampler_plot.get_complex_moduli(1)[0], dtype=complex)
    M_rep = np.asarray(model.ISD_matrix(z_rep, jnp.conj(z_rep)).real)
    M_inv_rep = np.linalg.inv(M_rep)

    # Gaussian samples
    eigvals_M, eigvecs_M = np.linalg.eigh(M_rep)
    sigma_sq = Nmax_plot / (n_fl * s_min_plot)
    scales = np.sqrt(sigma_sq * np.abs(eigvals_M))
    rng_plot = np.random.default_rng(0)
    N_samp = 50_000
    z_std = rng_plot.standard_normal((N_samp, n_fl))
    h_gauss = np.round((z_std * scales[None, :]) @ eigvecs_M.T).astype(int)
    h_gauss_norms = np.sqrt(np.sum(h_gauss**2, axis=1))
    h_gauss = h_gauss[h_gauss_norms > 0]
    h_gauss_norms = h_gauss_norms[h_gauss_norms > 0]

    # Tadpole of Gaussian samples
    hMh_gauss = np.abs(s_min_plot * np.einsum('hi,ij,hj->h', h_gauss.astype(float), M_inv_rep, h_gauss.astype(float)))

    # Uniform-from-box samples (for comparison)
    bf_plot = jvc.flux_bounding.bounded_fluxes(model, sampler=sampler_plot, Nmax=Nmax_plot)
    bf_plot.compute_eigenvalue_bounds(200, verbose=False)
    h_unif = bf_plot._sample_h_vectors(min(N_samp, 50_000), rng=np.random.default_rng(0))
    h_unif_norms = np.sqrt(np.sum(h_unif**2, axis=1))
    hMh_unif = np.abs(s_min_plot * np.einsum('hi,ij,hj->h', h_unif.astype(float), M_inv_rep, h_unif.astype(float)))

    # Dataset B (if available locally)
    if DATASET_B_PATH is not None:
        datsB = jvc.util.load_zipped_pickle(str(DATASET_B_PATH))
        h_B_norms = np.sqrt(np.sum(datsB[:, 12:18]**2, axis=1))
        has_datsB = True
    else:
        h_B_norms = np.array([])
        has_datsB = False

    fig, axes = plt.subplots(1, 3, figsize=(16, 4.5))

    # Panel 1: ||h|| distribution
    bins_h = np.arange(0, 60, 1)
    axes[0].hist(h_unif_norms, bins=bins_h, alpha=0.5, density=True, label='Uniform (box)', color='grey')
    axes[0].hist(h_gauss_norms, bins=bins_h, alpha=0.6, density=True, label=r'Gaussian $\mathcal{N}(0,\sigma^2 M)$', color='coral')
    if has_datsB:
        axes[0].hist(h_B_norms, bins=bins_h, alpha=0.5, density=True, label='Dataset B', color='steelblue')
    axes[0].set_xlabel(r'$\|h\|$')
    axes[0].set_ylabel('density')
    axes[0].set_title(r'$\|h\|$ distribution')
    axes[0].legend(fontsize=9)
    axes[0].set_xlim(0, 50)

    # Panel 2: Continuous tadpole distribution
    bins_t = np.linspace(0, 100, 50)
    axes[1].hist(hMh_unif, bins=bins_t, alpha=0.5, density=True, label='Uniform', color='grey')
    axes[1].hist(hMh_gauss, bins=bins_t, alpha=0.6, density=True, label='Gaussian', color='coral')
    axes[1].axvline(Nmax_plot, color='red', linestyle='--', linewidth=2, label=f'$N_{{\\max}}={Nmax_plot}$')
    axes[1].set_xlabel(r'$s_{\min} \cdot h^T M^{-1} h$')
    axes[1].set_ylabel('density')
    axes[1].set_title('Continuous tadpole')
    axes[1].legend(fontsize=9)
    axes[1].set_xlim(0, 100)

    # Panel 3: Tadpole pass rate vs Nmax
    Nmax_range = np.arange(1, 201)
    pass_gauss = [np.sum(hMh_gauss <= nm) / len(hMh_gauss) * 100 for nm in Nmax_range]
    pass_unif = [np.sum(hMh_unif <= nm) / len(hMh_unif) * 100 for nm in Nmax_range]
    axes[2].plot(Nmax_range, pass_gauss, color='coral', linewidth=2, label='Gaussian')
    axes[2].plot(Nmax_range, pass_unif, color='grey', linewidth=2, label='Uniform (box)')
    axes[2].set_xlabel(r'$N_{\max}$')
    axes[2].set_ylabel('% passing tadpole')
    axes[2].set_title('Tadpole acceptance rate')
    axes[2].legend()

    plt.suptitle(f'Gaussian prior vs uniform box sampling (Nmax={Nmax_plot})', y=1.02, fontsize=13)
    plt.tight_layout()
    plt.show()

    # Print key numbers
    print(f"Tadpole acceptance at Nmax={Nmax_plot}:")
    print(f"  Gaussian: {np.sum(hMh_gauss <= Nmax_plot)/len(hMh_gauss)*100:.1f}%")
    print(f"  Uniform:  {np.sum(hMh_unif <= Nmax_plot)/len(hMh_unif)*100:.1f}%")
    print(f"  → Gaussian is {np.sum(hMh_gauss<=Nmax_plot)/max(np.sum(hMh_unif<=Nmax_plot),1):.0f}× more efficient")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Legacy: uniform sampling of fluxes from a box#

Before ISD-based sampling, a simple baseline was to draw integer fluxes uniformly from a hypercube and Newton-refine every survivor. We include this baseline to make the efficiency gain from ISD-aware proposals visible, not because it is recommended for production scans.

if RUN_BENCHMARKS:
    # Legacy uniform-box sampling for comparison.
    kwargs_box = {"Q": model.Q(), "return_flag": True,
                  "constraints": None, "remove_NANs": True,
                  "in_axes": (0, 0, 0)}
    linearised_shifts_box = vmapping_func(
        model.linearised_shifts, mode="box", **kwargs_box)

    t0 = time.perf_counter()
    moduli_box, tau_box, fluxes_box, residuals_box = (
        model.sample_SUSY_flux_vacua(
            N=2 * 10**2,
            rns_key=rns_key,
            sampler=sampler,
            optimiser=linearised_shifts_box,
            objective_fct=DW_batch,
            vmap_dim=10**2,
            print_progress=True,
        )
    )
    dt = time.perf_counter() - t0
    print(f"Legacy uniform-box: {len(moduli_box)} vacua in {dt:.1f}s")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

ISD mode comparison#

sample_critical_points supports four ISD sampling modes for generating flux candidates. Their relative performance depends on the moduli region and on whether one is targeting SUSY vacua or broader scalar-potential critical points.

Mode

Input

Completes via

Best for

F

random \(f\)

\(h = M^{-1}\Sigma(f - c_0 h)\)

Mixed SUSY + non-SUSY candidates

H

random \(h\)

\(f = sM\Sigma h + c_0 h\)

SUSY-dominated candidates

ISD+

random \([f_2 \mid h_2]\)

compute \([f_1 \mid h_1]\) via \(\mathcal{N}\)

Often high SUSY yield

ISD-

random \([f_1 \mid h_1]\)

compute \([f_2 \mid h_2]\) via \(\mathcal{N}^{-1}\)

Useful for non-SUSY searches

The disabled benchmark cell compares all four modes on the same candidate budget.

if RUN_BENCHMARKS:
    modes = ["F", "H", "ISD+", "ISD-"]
    mode_results = {}

    for mode in modes:
        print(f"\n{'='*60}")
        print(f"  Mode: {mode}")
        print(f"{'='*60}")
        finder_m = model
        res = finder_m.sample_critical_points(Q=Nmax, 
                                              noscale=True, 
                                              sampler=sampler,         
                                              n_target=500,       # large target so we don't stop early
                                            n_batch=500,
                                            max_batches=1,      # single batch of 500 candidates
                                            isd_mode=mode,
                                            solver="scipy",
                                            verbose=True,
                                        )
        mode_results[mode] = res
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

if RUN_BENCHMARKS:
    # Summary table
    print(f"{'Mode':<8} {'Found':>6} {'SUSY':>6} {'non-SUSY':>9} {'minima':>7} {'saddle':>7} {'min rate':>9} {'non-SUSY%':>10}")
    print("-" * 72)
    for mode in modes:
        res = mode_results[mode]
        n = len(res)
        n_s = sum(1 for r in res if r.get('is_susy', False))
        n_ns = n - n_s
        n_m = sum(1 for r in res if r.get('is_minimum', False))
        n_sad = n - n_m
        mr = f"{100*n_m/n:.0f}%" if n > 0 else "—"
        nsr = f"{100*n_ns/n:.0f}%" if n > 0 else "—"
        print(f"{mode:<8} {n:>6} {n_s:>6} {n_ns:>9} {n_m:>7} {n_sad:>7} {mr:>9} {nsr:>10}")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Solver comparison: Newton vs scipy#

sample_critical_points supports two main local-solver backends:

  • Newton (solver="newton"): JIT-compiled Newton’s method using model.dV_x and model.ddV_x. Converges quadratically when close to a solution.

  • scipy (solver="scipy"): Uses scipy.optimize.root with the hybrid Powell method. It can be robust per candidate, but it is less naturally batched than JAX-native solvers.

The disabled benchmark cell compares them on the same ISD- flux candidates.

if RUN_BENCHMARKS:
    solvers = ["newton", "scipy"]
    solver_results = {}

    for slv in solvers:
        print(f"\n{'='*60}")
        print(f"  Solver: {slv}")
        print(f"{'='*60}")
        finder_s = model
        t_start = time.perf_counter()
        res = finder_s.sample_critical_points(
            Q=Nmax, 
            noscale=True, 
            sampler=sampler,
            n_target=500,
            n_batch=500,
            max_batches=1,
            isd_mode="ISD-",
            solver=slv,
            verbose=True,
        )
        t_end = time.perf_counter()
        solver_results[slv] = {'results': res, 'time': t_end - t_start}

    # Comparison
    print(f"\n{'Solver':<10} {'Found':>6} {'SUSY':>6} {'non-SUSY':>9} {'minima':>7} {'time (s)':>9}")
    print("-" * 52)
    for slv in solvers:
        res = solver_results[slv]['results']
        n = len(res)
        n_s = sum(1 for r in res if r.get('is_susy', False))
        n_ns = n - n_s
        n_m = sum(1 for r in res if r.get('is_minimum', False))
        t = solver_results[slv]['time']
        print(f"{slv:<10} {n:>6} {n_s:>6} {n_ns:>9} {n_m:>7} {t:>8.1f}s")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Advanced: performance tuning#

This section addresses performance and quality tuning for non-SUSY sampling at production scale. We compare the vectorised Adam and the hybrid solver against scipy/Newton, then walk through the seven ablation knobs from NB19 §10 (moduli_max, sampling region, ISD mode, solution magnitudes, \(N_{\max}\), and a production-tuned run).

Vectorised optax solver and hybrid approach#

The per-candidate loop in the Newton and scipy solvers becomes a bottleneck for large batches. sample_critical_points includes a vectorised optax Adam solver (solver="adam_v") that uses lax.scan for the step loop and vmap over the batch, compiling the warm-start into a single XLA kernel.

The hybrid solver (solver="hybrid") runs a short vectorised Adam warm-start to move candidates toward Newton’s convergence basin, then finishes with Newton’s quadratic convergence. In this model family it is often the best throughput/quality compromise, but the optimal settings remain model- and region-dependent.

Parameter analysis#

The key hyperparameters for the vectorised Adam warm-start are:

  • Learning rate schedule: cosine annealing with a moderate initial learning rate is a good starting point.

  • Gradient clipping: clip_by_global_norm(...) helps prevent runaway updates.

  • Step count: warm-starts should be long enough to enter a solver basin, but additional steps have diminishing returns.

  • Objective function: compare objectives on a small sweep before committing to a production run.

if RUN_BENCHMARKS:
    import optax

    # Generate a shared set of ISD- candidates for fair comparison
    finder_bench = model
    mod_pts = jnp.array(sampler.get_complex_moduli(500), dtype=complex)
    tau_pts = jnp.array(sampler.get_complex_tau(500), dtype=complex)
    x0_all, fl_all, _ = finder_bench._generate_flux_candidates(
        500, mod_pts, tau_pts, isd_mode="ISD-")
    n_bench = min(200, len(x0_all))
    print(f"Benchmarking on {n_bench} ISD- candidates\n")

    tol = 1e-8

    # --- Parameter sweep for Adam_v ---
    print("Learning rate schedule comparison (Adam, cosine annealing, 5000 steps):")
    print(f"{'lr':>6} {'conv':>6} {'min_res':>10} {'p10_res':>10} {'med_res':>10} {'time':>7}")
    print("-" * 55)
    for lr in [0.1, 0.3, 0.5, 0.7, 1.0]:
        sched = optax.cosine_decay_schedule(init_value=lr, decay_steps=5000)
        opt = optax.chain(optax.clip_by_global_norm(10.0), optax.adam(learning_rate=sched))
        t0 = time.perf_counter()
        _, res, conv = finder_bench._solve_dV_optax_batch(
            x0_all[:n_bench], fl_all[:n_bench], optimiser=opt,
            n_steps=5000, objective="dV2", tol=tol, sub_batch_size=50)
        dt = time.perf_counter() - t0
        sr = np.sort(res)
        print(f"{lr:>6.1f} {int(conv.sum()):>6} {sr[0]:>10.2e} {sr[int(0.1*len(sr))]:>10.2e} "
              f"{np.median(res):>10.2e} {dt:>6.1f}s")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

if RUN_BENCHMARKS:
    # --- Objective function comparison ---
    print("Objective function comparison (Adam cosine lr=0.5, 5000 steps):")
    print(f"{'objective':>10} {'conv':>6} {'min_res':>10} {'med_res':>10}")
    print("-" * 40)
    for obj in ["dV2", "log_dV2", "V"]:
        sched = optax.cosine_decay_schedule(init_value=0.5, decay_steps=5000)
        opt = optax.chain(optax.clip_by_global_norm(10.0), optax.adam(learning_rate=sched))
        _, res, conv = finder_bench._solve_dV_optax_batch(
            x0_all[:n_bench], fl_all[:n_bench], optimiser=opt,
            n_steps=5000, objective=obj, tol=tol, sub_batch_size=50)
        print(f"{obj:>10} {int(conv.sum()):>6} {np.min(res):>10.2e} {np.median(res):>10.2e}")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

if RUN_BENCHMARKS:
    # --- Full solver comparison: scipy vs Newton vs Adam_v vs Hybrid ---
    from scipy.optimize import root as scipy_root

    solver_bench = {}

    # 1. scipy
    t0 = time.perf_counter()
    scipy_conv = 0
    for j in range(n_bench):
        fl_np = fl_all[j].copy()
        def f(x, _fl=fl_np): return np.asarray(
            model.dV_x(jnp.array(x), jnp.array(_fl), noscale=True))
        def jac(x, _fl=fl_np): return np.asarray(
            model.ddV_x(jnp.array(x), jnp.array(_fl), noscale=True))
        r = scipy_root(f, x0=x0_all[j], method='hybr', jac=jac)
        if r.success and max(abs(r.fun)) < tol: scipy_conv += 1
    solver_bench['scipy'] = (scipy_conv, time.perf_counter() - t0)

    # 2. Newton
    t0 = time.perf_counter()
    newton_conv = 0
    for j in range(n_bench):
        _, _, c = finder_bench._solve_dV_newton_single(
            x0_all[j], fl_all[j], tol=tol, max_iters=300)
        if c: newton_conv += 1
    solver_bench['Newton'] = (newton_conv, time.perf_counter() - t0)

    # 3. Hybrid (Adam 1k + Newton)
    sched = optax.cosine_decay_schedule(init_value=0.5, decay_steps=1000)
    opt = optax.chain(optax.clip_by_global_norm(10.0), optax.adam(learning_rate=sched))
    t0 = time.perf_counter()
    x_warm, res_warm, _ = finder_bench._solve_dV_optax_batch(
        x0_all[:n_bench], fl_all[:n_bench], optimiser=opt,
        n_steps=1000, objective="dV2", tol=tol, sub_batch_size=50)
    hybrid_conv = int(np.sum(res_warm < tol))
    mask = (res_warm < 1.0) & (res_warm >= tol)
    for j in np.where(mask)[0]:
        _, res_n, conv_n = finder_bench._solve_dV_newton_single(
            x_warm[j], fl_all[j], tol=tol, max_iters=300)
        if conv_n: hybrid_conv += 1
    solver_bench['Hybrid 1k'] = (hybrid_conv, time.perf_counter() - t0)

    # Print comparison
    print(f"\n{'Solver':<14} {'Converged':>10} {'Rate':>8} {'Time':>8}")
    print("-" * 44)
    for name in ['scipy', 'Newton', 'Hybrid 1k']:
        nc, t = solver_bench[name]
        print(f"{name:<14} {nc:>10} {100*nc/n_bench:>7.1f}% {t:>7.1f}s")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Using the hybrid solver via sample_critical_points#

The hybrid solver is the recommended choice when you want to maximise the number of critical points found:

if RUN_BENCHMARKS:
    # The hybrid solver: vectorised Adam warm-start (1k steps) + Newton refinement
    finder_hybrid = model

    results_hybrid = finder_hybrid.sample_critical_points(
        Q=Nmax, 
        noscale=True, 
        sampler=sampler,
        n_target=50,
        n_batch=200,
        max_batches=3,
        isd_mode="ISD-",
        solver="hybrid",
        optax_steps=1000,           # Adam warm-start steps (1k is optimal)
        newton_max_iters=300,       # Newton refinement iterations
        verbose=True,
    )

    n_h = len(results_hybrid)
    n_ns_h = sum(1 for r in results_hybrid if not r.get('is_susy', False))
    n_min_h = sum(1 for r in results_hybrid if r.get('is_minimum', False))
    print(f"\nHybrid results: {n_h} critical points "
          f"({n_ns_h} non-SUSY, {n_min_h} minima)")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Performance analysis: moduli_max, sampling regions, and \(N_{\max}\)#

The moduli_max parameter, the moduli/dilaton sampling region, and \(N_{\max}\) all affect how many physical critical points are found versus how many candidate trajectories run away. This section is a recipe for designing a small sweep before a larger run.

Effect of moduli_max#

The moduli_max filter rejects solutions where any real coordinate component exceeds the threshold. Setting it too tight loses physical solutions; setting it too loose lets runaways pollute the benchmark.

if RUN_BENCHMARKS:
    # 10a. moduli_max sweep (ISD-, scipy, default region)
    print(f"{'moduli_max':>12} {'CritPts':>8} {'SUSY':>5} {'nonSUSY':>8} "
          f"{'minima':>7} {'saddle':>7} {'Time':>7}")
    print("-" * 62)

    sampler_default = jvc.data_sampler(model, moduli_bounds=(2., 5.),
        dilaton_bounds=(np.sqrt(3)/2, 10.), axion_bounds=(-0.5, 0.5), seed=42)

    for mmax in [10, 20, 50, 100, 200, 500, None]:
        f = model
        t0 = time.perf_counter()
        res = f.sample_critical_points(
            Q=Nmax, noscale=True, moduli_max=mmax, sampler=sampler_default, n_target=500, n_batch=500, max_batches=1,
            isd_mode="ISD-", solver="scipy", verbose=False)
        dt = time.perf_counter() - t0
        n_s = sum(1 for r in res if r.get('is_susy', False))
        n_m = sum(1 for r in res if r.get('is_minimum', False))
        label = str(mmax) if mmax is not None else "None"
        print(f"{label:>12} {len(res):>8} {n_s:>5} {len(res)-n_s:>8} "
              f"{n_m:>7} {len(res)-n_m:>7} {dt:>6.1f}s")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Effect of sampling region#

The moduli/dilaton sampling bounds determine where starting points are drawn. Tighter regions (closer to the origin) produce more physical solutions because the solver has less distance to travel.

if RUN_BENCHMARKS:
    # 10b. Sampling region sweep (ISD-, scipy, moduli_max=100)
    regions = [
        ((1., 2.), (np.sqrt(3)/2, 3.), "very tight: Im(z)∈(1,2), s∈(√3/2,3)"),
        ((1., 3.), (np.sqrt(3)/2, 5.), "tight: Im(z)∈(1,3), s∈(√3/2,5)"),
        ((2., 5.), (np.sqrt(3)/2, 10.), "default: Im(z)∈(2,5), s∈(√3/2,10)"),
        ((1., 10.), (np.sqrt(3)/2, 10.), "wide: Im(z)∈(1,10), s∈(√3/2,10)"),
        ((3., 8.), (1., 15.), "moderate: Im(z)∈(3,8), s∈(1,15)"),
        ((5., 15.), (2., 20.), "deep LCS: Im(z)∈(5,15), s∈(2,20)"),
    ]

    print(f"{'Region':>48} {'CritPts':>8} {'SUSY':>5} {'nonSUSY':>8} "
          f"{'minima':>7} {'Time':>7}")
    print("-" * 90)

    for mod_b, dil_b, label in regions:
        s = jvc.data_sampler(model, moduli_bounds=mod_b,
            dilaton_bounds=dil_b, axion_bounds=(-0.5, 0.5), seed=42)
        f = model
        t0 = time.perf_counter()
        res = f.sample_critical_points(
            Q=Nmax, noscale=True, moduli_max=100, sampler=s, n_target=500, n_batch=500, max_batches=1,
            isd_mode="ISD-", solver="scipy", verbose=False)
        dt = time.perf_counter() - t0
        n_s = sum(1 for r in res if r.get('is_susy', False))
        n_m = sum(1 for r in res if r.get('is_minimum', False))
        print(f"{label:>48} {len(res):>8} {n_s:>5} {len(res)-n_s:>8} "
              f"{n_m:>7} {dt:>6.1f}s")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

ISD mode comparison at tight region#

With a tighter sampling region and an explicit moduli_max filter, the relative performance of ISD modes can change. This is why the recommended production habit is to benchmark modes on the actual moduli patch of interest rather than transferring rankings from a different scan.

if RUN_BENCHMARKS:
    # 10c. ISD mode comparison at tight region, moduli_max=100
    sampler_tight = jvc.data_sampler(model, moduli_bounds=(1., 2.),
        dilaton_bounds=(np.sqrt(3)/2, 3.), axion_bounds=(-0.5, 0.5), seed=42)

    print(f"{'Mode':>6} {'Solver':>8} {'CritPts':>8} {'SUSY':>5} {'nonSUSY':>8} "
          f"{'minima':>7} {'Time':>7}")
    print("-" * 58)

    for mode in ['ISD-', 'F', 'H', 'ISD+']:
        for solver in ['scipy', 'hybrid']:
            f = model
            t0 = time.perf_counter()
            kw = {'optax_steps': 1000} if solver == 'hybrid' else {}
            res = f.sample_critical_points(
                Q=Nmax, noscale=True, moduli_max=100, sampler=sampler_tight, n_target=500, n_batch=500, max_batches=1,
                isd_mode=mode, solver=solver, verbose=False, **kw)
            dt = time.perf_counter() - t0
            n_s = sum(1 for r in res if r.get('is_susy', False))
            n_m = sum(1 for r in res if r.get('is_minimum', False))
            print(f"{mode:>6} {solver:>8} {len(res):>8} {n_s:>5} {len(res)-n_s:>8} "
                  f"{n_m:>7} {dt:>6.1f}s")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Solution magnitude distribution#

Without moduli_max, the vast majority of solutions are runaways. The tight region has the highest fraction of physical solutions.

if RUN_BENCHMARKS:
    # 10d. Solution magnitude distributions (ISD-, scipy, no moduli_max)
    dist_regions = [
        ((1., 2.), (np.sqrt(3)/2, 3.), "very tight"),
        ((2., 5.), (np.sqrt(3)/2, 10.), "default"),
        ((5., 15.), (2., 20.), "deep LCS"),
    ]

    for mod_b, dil_b, label in dist_regions:
        s = jvc.data_sampler(model, moduli_bounds=mod_b,
            dilaton_bounds=dil_b, axion_bounds=(-0.5, 0.5), seed=42)
        f = model
        res = f.sample_critical_points(
            Q=Nmax, noscale=True, moduli_max=None, sampler=s, n_target=500, n_batch=500, max_batches=1,
            isd_mode="ISD-", solver="scipy", verbose=False)

        if not res:
            print(f"{label}: 0 solutions")
            continue

        max_abs = np.array([float(np.max(np.abs(np.append(r['moduli'], r['tau']))))
                            for r in res])
        print(f"{label} ({len(res)} pts):")
        print(f"  max(|z|,|tau|): min={max_abs.min():.1f}, "
              f"median={np.median(max_abs):.1e}, max={max_abs.max():.1e}")
        for hi in [5, 10, 50, 100]:
            n_in = int(np.sum(max_abs <= hi))
            print(f"    ≤{hi:>5}: {n_in:>4}/{len(res)} ({100*n_in/len(res):>5.1f}%)")
        print()
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

\(N_{\max}\) sensitivity#

The tadpole constraint controls flux magnitudes. Larger \(N_{\max}\) allows more diverse fluxes but also more runaways.

if RUN_BENCHMARKS:
    # 10e. Nmax sweep (tight region, ISD-, scipy, moduli_max=100)
    print(f"{'Nmax':>6} {'CritPts':>8} {'SUSY':>5} {'nonSUSY':>8} "
          f"{'minima':>7} {'Time':>7}")
    print("-" * 48)

    for nmax in [10, 20, 50, 100, 200, 500]:
        f = model
        t0 = time.perf_counter()
        res = f.sample_critical_points(
            Q=nmax, noscale=True, moduli_max=100, sampler=sampler_tight, n_target=500, n_batch=500, max_batches=1,
            isd_mode="ISD-", solver="scipy", verbose=False)
        dt = time.perf_counter() - t0
        n_s = sum(1 for r in res if r.get('is_susy', False))
        n_m = sum(1 for r in res if r.get('is_minimum', False))
        print(f"{nmax:>6} {len(res):>8} {n_s:>5} {len(res)-n_s:>8} "
              f"{n_m:>7} {dt:>6.1f}s")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Production-style run: tight region + hybrid + moduli_max#

This disabled cell combines the settings selected by the earlier sweeps: a tight sampling region, the hybrid solver, and an explicit moduli_max cutoff. Treat it as a template for a production-style local run, not as a default tutorial cell.

if RUN_BENCHMARKS:
    # 10f. Production run: tight region + hybrid + moduli_max=100
    finder_prod = model

    results_prod = finder_prod.sample_critical_points(
        Q=Nmax, noscale=True, moduli_max=100, sampler=sampler_tight, n_target=200, n_batch=500, max_batches=5,
        isd_mode="ISD-", solver="hybrid", optax_steps=1000, verbose=True)

    n_p = len(results_prod)
    if n_p > 0:
        max_abs = np.array([float(np.max(np.abs(np.append(r['moduli'], r['tau']))))
                            for r in results_prod])
        Vs = np.array([r['V'] for r in results_prod])
        DWs = np.array([r['|DW|'] for r in results_prod])
        Nfluxes = np.array([r['Nflux'] for r in results_prod])

        print(f"\n{n_p} critical points found:")
        print(f"  max(|z|,|tau|): min={max_abs.min():.1f}, "
              f"median={np.median(max_abs):.1f}, max={max_abs.max():.1f}")
        print(f"  V: min={Vs.min():.2e}, median={np.median(Vs):.2e}")
        print(f"  |DW|: min={DWs.min():.2e}, median={np.median(DWs):.2e}")
        print(f"  Nflux: min={Nfluxes.min():.0f}, median={np.median(Nfluxes):.0f}, "
              f"max={Nfluxes.max():.0f}")

        # All solutions are within moduli_max
        print(f"\n  100% of solutions have max(|z|,|tau|) ≤ {max_abs.max():.1f} "
              f"(moduli_max={finder_prod.moduli_max})")
else:
    print("RUN_BENCHMARKS = False — skipping this benchmark cell.")

Summary of performance analysis#

Model-specific lessons for the h12=2 degree-18 fixture:

  1. moduli_max is a quality filter. Without it, many converged numerical critical points may be far outside the intended patch.

  2. Sampling region matters. Starting closer to the target patch often improves physical yield more than simply increasing the candidate budget.

  3. ISD-mode rankings are context-dependent. A mode that produces many raw candidates in one region may be less useful after patch and runaway filters are applied.

  4. Large \(N_{\max}\) is not automatically better. Higher flux budget can increase diversity, but it also increases the chance of expensive low-yield scans.

  5. Use the benchmark cells as a local calibration tool. Run short sweeps, inspect yield and residuals, and only then scale up n_batch, max_batches, or the moduli-region coverage.

Take-aways#

  • Benchmarking is a separate task from learning the sampling API; keep NB07 for workflow mechanics and use this notebook when choosing methods.

  • Bounded enumeration is complete only when the eigenvalue box is small enough; otherwise use stochastic bounded sampling or ISD-biased candidate generation.

  • Solver comparisons are sensitive to sampling region, moduli_max, \(N_{\max}\), and ISD mode; run small sweeps before launching production jobs.

  • Keep heavy cells opt-in so the tutorial suite remains laptop-feasible.

Further reading#