4D round q-Gaussian generator

examples/particles_generation/008_generate_q_gaussian.py

@@ -0,0 +1,99 @@
+# copyright ############################### #
+# This file is part of the Xpart Package.   #
+# Copyright (c) CERN, 2021.                 #
+# ######################################### #
+
+import json
+import numpy as np
+from matplotlib import pyplot as plt
+import xpart as xp
+import xtrack as xt
+
+def q_gaussian_1d(x, q, beta, normalize=False):
+    """
+    Args:
+        x:
+        q: q-parameter
+        beta: beta for q-Gaussian
+        normalize: if normalize area to 1
+
+    Returns:
+        q-Gaussian function defined on x
+
+    """
+    assert q < 5/3, "q must be less than 5/3"
+    arg = 1 - (1 - q) * beta * x**2
+    f = np.where(arg > 0, arg**(1 / (1 - q)), 0)
+    if normalize:
+        dx = x[1] - x[0]
+        area = np.sum(f) * dx
+        f /= area
+    return f
+
+bunch_intensity = 1e11
+sigma_z = 22.5e-2
+n_part = int(5e6)
+nemitt_x = 2e-6
+nemitt_y = 2.5e-6
+
+filename = ('../../../xtrack/test_data/sps_w_spacecharge'
+            '/line_no_spacecharge_and_particle.json')
+with open(filename, 'r') as fid:
+    ddd = json.load(fid)
+line = xt.Line.from_dict(ddd['line'])
+line.particle_ref = xp.Particles.from_dict(ddd['particle'])
+
+line.build_tracker()
+
+q = 1.3
+beta = 1
+
+x_norm, px_norm, y_norm, py_norm = xp.generate_round_4D_q_gaussian_normalised(
+    q=q,
+    beta=beta,
+    n_part=int(1e6)
+)
+
+particles = line.build_particles(
+                               zeta=0, delta=1e-3,
+                               x_norm=x_norm, # in sigmas
+                               px_norm=px_norm, # in sigmas
+                               y_norm=y_norm,
+                               py_norm=py_norm,
+                               nemitt_x=3e-6, nemitt_y=3e-6)
+
+# plot 1D q-gaussian against projection
+x = np.linspace(-10, 10, 1000)
+f = q_gaussian_1d(x=x, q=q, beta=beta, normalize=True)
+
+plt.close('all')
+fig1 = plt.figure(1, figsize=(6.4, 7))
+ax21 = fig1.add_subplot(3,1,1)
+ax22 = fig1.add_subplot(3,1,2)
+ax23 = fig1.add_subplot(3,1,3)
+ax21.plot(particles.x*1000, particles.px, '.', markersize=1)
+ax21.set_xlabel(r'x [mm]')
+ax21.set_ylabel(r'px [-]')
+ax22.plot(particles.y*1000, particles.py, '.', markersize=1)
+ax22.set_xlabel(r'y [mm]')
+ax22.set_ylabel(r'py [-]')
+ax23.plot(x, f, color='k', label=f'1D q-Gaussian q={q}, beta={beta}')
+ax23.hist(x_norm, bins=200, density=True, alpha=0.5,
+          label=f'normalised sampled q-Gaussian q={q}, beta={beta}, $x$ projection')
+ax23.hist(y_norm, bins=200, density=True, alpha=0.5,
+          label=f'normalised sampled q-Gaussian q={q}, beta={beta} $y$ projection')
+ax23.set_xlabel(r'normalised x')
+ax23.set_xlabel(r'normalised amplitude')
+ax23.legend()
+
+fig1.subplots_adjust(bottom=.08, top=.93, hspace=.33, left=.18,
+                     right=.96, wspace=.33)
+plt.show()
+
+
+
+
+
+
+
+

tests/test_transverse_q_gaussian_4d.py

@@ -0,0 +1,66 @@
+import numpy as np
+from scipy.optimize import curve_fit
+
+from xpart.transverse_generators.q_gaussian_round import (
+    generate_round_4D_q_gaussian_normalised)
+
+def test_transverse_q_gaussian_4d_sample_std():
+    """
+    test standard deviation is correct
+    """
+    q = 1.4
+    beta = 1
+    n_part = 10000000
+    x, _, _, _ = generate_round_4D_q_gaussian_normalised(q, beta, n_part)
+
+    # variance a function of q and beta
+    sample_variance = np.std(x)
+    expected_variance = np.sqrt(1 / (beta * (5 - 3 * q)))
+
+    assert np.isclose(sample_variance, expected_variance, atol=0.05), \
+        f"Sample variance {sample_variance} deviates from expected {expected_variance}"
+
+def q_gaussian_1d(x, q, beta, normalize=False):
+    """
+    Args:
+        x:
+        q: q-parameter
+        beta: beta for q-Gaussian
+        normalize: if normalize area to 1
+
+    Returns:
+        q-Gaussian function defined on x
+
+    """
+    assert q < 5/3, "q must be less than 5/3"
+    arg = 1 - (1 - q) * beta * x**2
+    f = np.where(arg > 0, arg**(1 / (1 - q)), 0)
+    if normalize:
+        dx = x[1] - x[0]
+        area = np.sum(f) * dx
+        f /= area
+    return f
+
+def test_transverse_q_gaussian_4d_sampler_returns_q0():
+    """
+    test that required q matches fitted q from scipy.optimize.curve fit
+    """
+    q = 1.4
+    beta = 1
+    n_part = 10000000
+    x, _, _, _ = generate_round_4D_q_gaussian_normalised(q, beta, n_part)
+    # Generate histogram
+    bins = np.linspace(-10, 10, 1000)  # 1000 bins between -10 and 10
+    counts, bin_edges = np.histogram(x, bins=bins, density=True)  # density=True to normalize to PDF
+
+    # Calculate bin centers for fitting
+    bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
+    popt, pcov = curve_fit(q_gaussian_1d, bin_centers, counts, p0=[1.5, 1, 1.0])
+
+    assert np.isclose(popt[0], q, atol=0.05), \
+        f"Fitted q from samples {popt[0]} does not match requested q {q} "
+
+
+# test that variable F(q, beta) works?
+
+

xpart/__init__.py

@@ -22,6 +22,7 @@
 from .transverse_generators import generate_2D_gaussian
 from .transverse_generators import (generate_hypersphere_2D, generate_hypersphere_4D,
                                     generate_hypersphere_6D)
+from .transverse_generators import generate_round_4D_q_gaussian_normalised
 
 from .longitudinal import generate_longitudinal_coordinates
 from .longitudinal.generate_longitudinal import _characterize_line

xpart/transverse_generators/__init__.py

@@ -8,3 +8,4 @@
 from .pencil import generate_2D_pencil, generate_2D_pencil_with_absolute_cut
 from .gaussian import generate_2D_gaussian
 from .hypersphere import generate_hypersphere_2D, generate_hypersphere_4D, generate_hypersphere_6D
+from .q_gaussian_round import generate_round_4D_q_gaussian_normalised

xpart/transverse_generators/q_gaussian_round.py

@@ -0,0 +1,151 @@
+#################################################
+# This code randomly samples 4D                 #
+# q-Gaussian distributions (1<q<5/3) using the  #
+# sampling methods of Batygin                   #
+# https://doi.org/10.1016/j.nima.2004.10.029    #
+# and the 4D q-Gaussian formula derived in      #
+# https://cds.cern.ch/record/2912366?ln=en      #
+# These distributions are non-factorizable.     #
+#################################################
+
+import numpy as np
+from scipy.special import gamma
+from scipy.interpolate import interp1d
+
+
+def generate_radial_distribution(q, beta, F):
+    """
+    Compute the 4D radial distribution function for a round q-Gaussian.
+    This can be numerically unstable if extreme values of q, beta as the
+    sample space needs to be increased (eg for very high q or beta)
+
+    Parameters:
+        q (float): q-parameter (1 < q < 5/3).
+        beta (float): Scale parameter, (beta > 1)
+        F: sample space parameter, can be user defined
+
+    Returns:
+        tuple: (f_F, F) where f_F is the radial distribution, and F is the radial coordinate array.
+    """
+    assert q > 1, "q must be greater than 1"
+    assert q < 5/3, "q must be less than 5/3"
+    assert beta > 0, "beta must be greater than 0"
+
+    term1 = -(beta**2) * (q - 3) * (q**2 - 1) / 4 / np.pi**2
+    if abs(q - 1) < 1e-2:
+        term2 = -1 / (1 - q)
+    else:
+        term2 = gamma(q / (q - 1)) / gamma(1 / (q - 1))
+
+    term3 = (1 + beta * (q - 1) * F) ** (1 / (1 - q) - 3 / 2)
+    return term1 * term2 * term3, F
+
+
+def generate_pdf(f_F, F):
+    """
+    Compute the PDF g(F) from f(F) using the Abel transform in 4D.
+    Cleans up the boundaries and gives correct normalisation factor.
+
+    Parameters:
+        f_F (np.ndarray): Distribution array.
+        F (np.ndarray): Radial coordinate array.
+
+    Returns:
+        np.ndarray: Transformed PDF g(F).
+    """
+    f_F = f_F.copy()
+    f_F[0] = f_F[-1] = 0  # Boundary cleanup
+    return np.pi**2 * f_F * F
+
+
+def generate_cdf(g_F, F):
+    """
+    Compute the cumulative distribution function (CDF) of g(F).
+
+    Parameters:
+        g_F (np.ndarray): PDF values.
+        F (np.ndarray): Radial coordinates.
+
+    Returns:
+        np.ndarray: CDF of g(F).
+    """
+    delta_F = np.diff(F, prepend=0)
+    cdf = np.cumsum(g_F * delta_F) # todo: rewrite with scipy.integrate.quad(g_F, -np.inf, np.inf)
+    cdf = np.clip(cdf, 0, np.inf)
+    return cdf
+
+
+def sample_from_inv_cdf(Np, cdf_g, F):
+    """
+    Sample F values from the inverse CDF of g(F).
+
+    Parameters:
+        Np (int): Number of particles to sample.
+        cdf_g (np.ndarray): CDF of g(F).
+        F (np.ndarray): Original F grid.
+
+    Returns:
+        np.ndarray: Sampled F values (F_G).
+    """
+    cdf_g /= cdf_g[-1]  # normalize
+    uniform_samples = np.random.uniform(0, 1, Np)
+    interpolator = interp1d(
+        cdf_g, F, kind="nearest", bounds_error=False, fill_value=(F[0], F[-1])
+    )
+    return interpolator(uniform_samples)
+
+
+def generate_random_a(F_G):
+    """
+    Generate A_x and A_y coordinates based on F_G distribution.
+
+    Parameters:
+        F_G (np.ndarray): Sampled F values.
+
+    Returns:
+        tuple: (A_x, A_y) arrays.
+    """
+    A_X_SQ = np.random.uniform(0, F_G)
+    A_x = np.sqrt(A_X_SQ)
+    A_y = np.sqrt(F_G - A_X_SQ)
+    return A_x, A_y
+
+
+def generate_round_4D_q_gaussian_normalised(q, beta, n_part, sample_space=None):
+    """
+    Generate particles sampled from a 4D round q-Gaussian distribution.
+
+    Parameters:
+        q (float): q-Gaussian shape parameter. Must satisfy 1 < q < 5/3.
+        beta (float): Scale parameter.
+        n_part (int): Number of particles to generate.
+        sample_space (array-like, optional): 1D array of radius values used for sampling.
+            Defaults to np.linspace(0, 3000, 100000) if None.
+
+    Returns:
+        tuple: Arrays of normalised transverse coordinates (x, px, y, py).
+    """
+    if sample_space is None:
+        F = np.linspace(0, 3e2, 1e6)
+    else:
+        F = sample_space
+
+    f_F, F = generate_radial_distribution(q, beta, F)  # 4D distribution
+    g_F = generate_pdf(f_F, F)  # PDF of 4D distribution
+    cdf_g = generate_cdf(g_F, F)  # CDF
+    F_G = sample_from_inv_cdf(n_part, cdf_g, F)  # Inverse function
+    A_x, A_y = generate_random_a(F_G)  # random generator distributed like F_G
+
+    # Sample angles for all particles
+    beta_x = np.random.uniform(0, 2 * np.pi, n_part)
+    beta_y = np.random.uniform(0, 2 * np.pi, n_part)
+
+    # Compute positions and momenta
+    x = A_x * np.cos(beta_x)
+    px = -A_x * np.sin(beta_x)
+    y = -A_y * np.cos(beta_y)
+    py = A_y * np.sin(beta_y)
+
+    return x, px, y, py
+
+