jaxvacua quickstart

jaxvacua quickstart#

Find your first SUSY flux vacuum end-to-end in under 10 minutes on a CPU.

What this notebook does:

Loads a small bundled Calabi–Yau model (no HuggingFace download, no CYTools).

Samples ~200 random initial guesses for the moduli, axio-dilaton and integer flux quanta.

Newton-solves the F-term equations \(D_I W = 0\) to find SUSY vacua.

Inspects the converged set — counts, \(|W_0|\) distribution, \(g_s = 1/\mathrm{Im}\,\tau\).

What this notebook deliberately skips (covered in complete_runthrough_*.ipynb):

The HuggingFace-hosted CY database (LCSDatabase).

The vacua vault (storing / designating / pushing solutions).

Mass spectra, stability checks, non-SUSY critical points.

Runs on Colab CPU; no GPU / Metal needed.

Setup#

# Colab-friendly install + JAX float64 setup. Idempotent.
try:
    import jaxvacua
except ImportError:
    %pip install --quiet jaxvacua
    import jaxvacua

import jax
jax.config.update("jax_enable_x64", True)

import jaxvacua as jvc
import jax.numpy as jnp
import numpy as np
from scipy.optimize import root
import matplotlib.pyplot as plt

# On Apple Silicon, jaxvacua auto-falls back from Metal to CPU
# (Metal does not yet support complex numbers).
print(f"jax backend: {jax.default_backend()}")
print(f"jaxvacua precision: {jvc.precision}")

1. Load a model#

We use a small bundled Kreuzer–Skarke model with \(h^{1,2}=2\) (two complex-structure moduli). maximum_degree=0 disables instanton corrections — the prepotential stays purely polynomial, which keeps everything fast and analytic.

model = jvc.FluxVacuaFinder(
    h12=2, model_ID=1, model_type="KS",
    maximum_degree=0,    # polynomial prepotential, no instantons
)
print(f"h^(1,2) = {model.h12}     (number of complex-structure moduli)")
print(f"chi      = {int(model.lcs_tree.chi)}    (Euler characteristic)")
print(f"intnums shape: {model.lcs_tree.intnums.shape}")

2. Sample initial guesses#

data_sampler draws N random points \((z, \tau, \text{fluxes})\) inside user-defined bounds. We use the simplest path: uniform sampling, integer fluxes in \([-5, 5]\), moduli filtered to lie inside the Kähler cone.

sampler = jvc.data_sampler(
    model,
    moduli_bounds=(2., 5.),
    dilaton_bounds=(2., 10.),
    axion_bounds=(-0.5, 0.5),
    flux_bounds=(-5, 5),
    seed=42,
)
z0, tau0, fluxes0 = sampler.initial_guesses(
    N=200, filter_moduli=True, include_fluxes=True,
)
print(f"sampled {len(z0)} candidate guesses")
print(f"  z0     shape: {z0.shape}")
print(f"  tau0   shape: {tau0.shape}")
print(f"  fluxes shape: {fluxes0.shape}")

3. Refine the F-term equations#

The quickstart uses SciPy’s root(..., method="hybr") on the real-vector form of the F-term system. This is a robust small-batch interface because we can pass the analytic Jacobian model.dDW_x directly.

JAXVacua also provides its own damped Newton iterator, model.newton_method_flux_vacua, which is useful in batched workflows and appears in the dedicated vacuum-finding notebooks. We show both interfaces here: SciPy is used for the main scan, and the internal Newton method is demonstrated on one candidate below.

def solve_from_guesses(model, z0, tau0, fluxes0, tol=1e-10, residual_tol=1e-6):
    """Solve DW=0 with scipy.root from each guess."""
    moduli_out, tau_out, flux_out, residuals = [], [], [], []
    for i in range(len(z0)):
        x0 = model._convert_complex_to_real(
            z0[i], jnp.conj(z0[i]), tau0[i], jnp.conj(tau0[i])
        )
        r = root(model.DW_x, x0, args=(fluxes0[i],),
                 jac=model.dDW_x, method="hybr", tol=tol)
        residual = float(np.max(np.abs(r.fun)))
        if residual < residual_tol:
            m, _, t, _ = model._convert_real_to_complex(r.x)
            if not np.any(np.isnan(np.append(m, t))):
                moduli_out.append(np.asarray(m))
                tau_out.append(complex(t))
                flux_out.append(np.asarray(fluxes0[i]))
                residuals.append(residual)
    return moduli_out, tau_out, flux_out, residuals

moduli, taus, fluxes, residuals = solve_from_guesses(
    model, z0[:50], tau0[:50], fluxes0[:50],
)
print(f"converged: {len(moduli)} / 50  ({100*len(moduli)/50:.0f}%)")

Optional: internal Newton iterator#

For comparison, the same FluxVacuaFinder object can take damped Newton steps internally. The call below is deliberately just a single-candidate demonstration; the SciPy loop above remains the clearer quickstart route for collecting a small table of vacua.

z_newton, tau_newton, res_newton = model.newton_method_flux_vacua(
    z0[0], tau0[0], fluxes0[0],
    mode="SUSY",
    solver_mode="real",
    step_size_Newton=0.2,
    tol=1e-8,
    max_iters=25,
)

DW_newton = model.DW(
    z_newton, jnp.conj(z_newton),
    tau_newton, jnp.conj(tau_newton),
    fluxes0[0],
)
print(f"internal Newton residual reported: {float(res_newton):.2e}")
print(f"direct max |DW| check:          {float(jnp.max(jnp.abs(DW_newton))):.2e}")

4. Inspect the vacua#

For each converged vacuum we compute:

\(|W_0| = |\langle W \rangle|\), the gauge-invariant value of the superpotential — small \(|W_0|\) vacua are physically interesting (KKLT-style).
\(g_s = 1/\mathrm{Im}\,\tau\), the string coupling.
the residual \(\max|D_I W|\) as a sanity check (should be tiny).

import pandas as pd
rows = []
for m, t, f, r in zip(moduli, taus, fluxes, residuals):
    W0 = abs(complex(model.superpotential_gauge_invariant(
        jnp.array(m), complex(t), jnp.array(f)
    )))
    rows.append({
        "|W0|": W0,
        "gs":   1. / float(np.imag(t)),
        "residual": r,
        "Re(z1)": float(np.real(m[0])),
        "Im(z1)": float(np.imag(m[0])),
    })
df = pd.DataFrame(rows).sort_values("|W0|").reset_index(drop=True)
print(f"converged vacua: {len(df)}")
print(f"  |W0|     min={df['|W0|'].min():.3e}  median={df['|W0|'].median():.3e}")
print(f"  gs       min={df['gs'].min():.3f}     max={df['gs'].max():.3f}")
print(f"  residual max={df['residual'].max():.2e}")
df.head(5)

fig, ax = plt.subplots(figsize=(6, 4))
sc = ax.scatter(df["Re(z1)"], df["Im(z1)"], c=np.log10(df["|W0|"]),
                s=30, cmap="viridis", edgecolor="k", linewidth=0.3)
ax.set_xlabel(r"$\mathrm{Re}\,z^1$")
ax.set_ylabel(r"$\mathrm{Im}\,z^1$")
ax.set_title(f"{len(df)} SUSY flux vacua — coloured by $\\log_{{10}}|W_0|$")
fig.colorbar(sc, ax=ax, label=r"$\log_{10}|W_0|$")
fig.tight_layout()
plt.show()

Take-aways#

JAXVacua’s end-to-end SUSY workflow fits in four steps: load model → sample initial guesses → Newton-refine to \(D_I W = 0\) → inspect.
The FluxVacuaFinder class bundles all four — jvc.data_sampler for draws, model.newton_method_flux_vacua(..., mode="SUSY") (or the manual scipy.optimize.root route used above) for refinement, and model.W0/model.DW for verification.
The whole pipeline is JIT-compiled and vmap-friendly. The cells above sample ~200 initial guesses in seconds on a laptop CPU.
For production runs (1k-1M vacua, ensembles across geometries), graduate to the workflows in NB07 (ISD sampling) and NB08 (flux bounding); the API is identical.