Neural importance sampling with MadNIS

The integrator in the previous example samples uniformly and reweights, which is why its Monte Carlo error only shrinks as the square root of the sample count. MadNIS is a separate package that trains a normalizing flow to sample close to the integrand itself, cutting that error down for the same number of evaluations. This example wraps the previous integrator’s mapping and cross section as a MadNIS integrand.

Setting up the physics

Everything up to the integrand itself is the same as in the previous example, except that the random numbers are now PyTorch tensors, since MadNIS trains with PyTorch:

import os

import lhapdf  # only used below to locate the installed PDF set
import torch
import madspace as ms
from madnis.integrator import Integrator

torch.set_default_dtype(torch.float64)

E_CM = 13000.0
masses = [0.0, 0.0, 173.0, 173.0, 0.0]  # g g -> t t~ g
pids = [21, 21, 6, -6, 21]

O = ms.Observable
cuts = ms.Cuts([
    ms.CutItem(O(pids, O.obs_pt, [O.jet_pids]), min=20.0),
    ms.CutItem(O(pids, O.obs_eta_abs, [O.jet_pids]), max=5.0),
    ms.CutItem(O(pids, O.obs_delta_r, [O.jet_pids]), min=0.4),
])
mapping = ms.PhaseSpaceMapping(masses, E_CM, mode="propagator", cuts=cuts)

proc_dir = "PROC_ggttg"
me_path = os.path.join(proc_dir, "lib", "libmadmatrix_P0_gg_ttxg_scalar.so")
param_card = os.path.join(proc_dir, "Cards", "param_card.dat")

ctx = ms.default_context()
api = ctx.load_matrix_element(me_path, param_card)
MEI = ms.MatrixElement.MatrixElementInput
MEO = ms.MatrixElement.MatrixElementOutput
matrix_element = ms.MatrixElement(
    api, [MEI.momenta_in, MEI.alpha_s_in], [MEO.matrix_element_out], True
)

PDF_SET = "NNPDF40_lo_as_01180"
pdf_dir = os.path.join(lhapdf.paths()[0], PDF_SET)
pdf_grid = ms.PdfGrid(os.path.join(pdf_dir, f"{PDF_SET}_0000.dat"))
pdf_grid.initialize_globals(ctx)
alphas_grid = ms.AlphaSGrid(os.path.join(pdf_dir, f"{PDF_SET}.info"))
alphas_grid.initialize_globals(ctx)
alpha_s = ms.RunningCoupling(alphas_grid)

scale = ms.EnergyScale(
    particle_count=5, type=ms.EnergyScale.DynamicalScaleType.half_transverse_mass
)
cross_section = ms.DifferentialCrossSection(
    matrix_element=matrix_element,
    cm_energy=E_CM,
    running_coupling=alpha_s,
    energy_scale=scale,
    pid_options=[[21, 21]],
    pdf1=pdf_grid, pdf2=pdf_grid,
    input_momentum_fraction=True,
)

Writing the integrand

MadNIS draws unit-hypercube points and asks for the integrand at each of them. The mapping turns those points into momenta, and its Jacobian, det, has to be folded into the returned value along with the cross section, exactly as in the manual integrator:

def integrand(x):
    momenta, x1, x2, det = mapping.map_forward([x])
    pdf_id = torch.zeros(x.shape[0], dtype=torch.int32)
    return cross_section(momenta, x1, x2, pdf_id) * det

Training and integrating

madnis.integrator.Integrator takes the integrand and its dimension. Training adapts the flow to the integrand; integrate() then draws fresh samples from the trained flow to compute the final estimate:

integrator = Integrator(integrand, dims=mapping.random_dim())
integrator.train(200)
result, error = integrator.integrate(200000)
print(f"sigma = {result:.2f} +- {error:.2f} pb  (rel. error {error / result * 100:.2f}%)")

The same number of samples now gives a much smaller error than the non-adaptive integrator:

sigma = 239.64 +- 0.34 pb  (rel. error 0.14%)