diff --git a/examples/particles_generation/008_generate_q_gaussian.py b/examples/particles_generation/008_generate_q_gaussian.py new file mode 100644 index 00000000..0aa9aff8 --- /dev/null +++ b/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() + + + + + + + + diff --git a/tests/test_transverse_q_gaussian_4d.py b/tests/test_transverse_q_gaussian_4d.py new file mode 100644 index 00000000..4f05a432 --- /dev/null +++ b/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? + + diff --git a/xpart/__init__.py b/xpart/__init__.py index 5695e575..8609d0ea 100644 --- a/xpart/__init__.py +++ b/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 diff --git a/xpart/transverse_generators/__init__.py b/xpart/transverse_generators/__init__.py index d1d6e467..7740ee8a 100644 --- a/xpart/transverse_generators/__init__.py +++ b/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 diff --git a/xpart/transverse_generators/q_gaussian_round.py b/xpart/transverse_generators/q_gaussian_round.py new file mode 100644 index 00000000..ce34e578 --- /dev/null +++ b/xpart/transverse_generators/q_gaussian_round.py @@ -0,0 +1,151 @@ +################################################# +# This code randomly samples 4D # +# q-Gaussian distributions (1 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 + +