Hypersphere distribution and nemitt_zeta in build_particle

tests/test_hypersphere.py

@@ -0,0 +1,81 @@
+# copyright ############################### #
+# This file is part of the Xpart Package.   #
+# Copyright (c) CERN, 2021.                 #
+# ######################################### #
+
+import numpy as np
+from scipy import stats
+
+import xpart as xp
+
+
+def test_hypersphere_2D():
+    num_particles = int(20e4)
+    r = 2.5
+
+    means = []
+    stds  = []
+    for seed in [1,2,3]:
+        x_norm, px_norm = xp.hypersphere_2D(num_particles,r = r, rng_seed = seed)
+
+        data = np.array([x_norm,px_norm]).T
+        bins = [np.linspace(-r,r,50) for i in range(2)]
+
+        bincounts = stats.binned_statistic_dd(data, 1,bins=bins,statistic='count').statistic
+        bincounts = bincounts[bincounts!=0]
+
+        means.append(np.mean(bincounts))
+        stds.append(np.std(bincounts))
+
+
+    assert np.allclose(np.diff(means)/np.mean(means),0,rtol=0, atol=1e-2)
+    assert np.allclose(np.array(stds)/np.array(means),0,rtol=0, atol=0.5)
+
+
+def test_hypersphere_4D():
+    num_particles = int(20e4)
+    rx = 2.5
+    ry = 4.2
+
+    means = []
+    stds  = []
+    for seed in [1,2,3]:
+        x_norm, px_norm, y_norm, py_norm = xp.hypersphere_4D(num_particles,rx=rx,ry=ry, rng_seed = seed)
+
+        data = np.array([x_norm,px_norm,y_norm,py_norm]).T
+        bins = [np.linspace(-rx,rx,50) for i in range(2)] + [np.linspace(-ry,ry,50) for i in range(2)]
+
+        bincounts = stats.binned_statistic_dd(data, 1,bins=bins,statistic='count').statistic
+        bincounts = bincounts[bincounts!=0]
+
+        means.append(np.mean(bincounts))
+        stds.append(np.std(bincounts))
+
+
+    assert np.allclose(np.diff(means)/np.mean(means),0,rtol=0, atol=1e-2)
+    assert np.allclose(np.array(stds)/np.array(means),0,rtol=0, atol=0.5)
+
+
+def test_hypersphere_6D():
+    num_particles = int(20e4)
+    rx = 2.5
+    ry = 4.2
+    rzeta = 1.2
+
+    means = []
+    stds  = []
+    for seed in [1,2,3]:
+        x_norm, px_norm, y_norm, py_norm, zeta_norm, pzeta_norm = xp.hypersphere_6D(num_particles,rx=rx,ry=ry,rzeta=rzeta, rng_seed = seed)
+
+        data = np.array([x_norm,px_norm,y_norm,py_norm,zeta_norm, pzeta_norm]).T
+        bins = [np.linspace(-rx,rx,20) for i in range(2)] + [np.linspace(-ry,ry,20) for i in range(2)] + [np.linspace(-rzeta,rzeta,20) for i in range(2)]
+
+        bincounts = stats.binned_statistic_dd(data, 1,bins=bins,statistic='count').statistic
+        bincounts = bincounts[bincounts!=0]
+
+        means.append(np.mean(bincounts))
+        stds.append(np.std(bincounts))
+
+
+    assert np.allclose(np.diff(means)/np.mean(means),0,rtol=0, atol=1e-2)
+    assert np.allclose(np.array(stds)/np.array(means),0,rtol=0, atol=0.5)

xpart/__init__.py

@@ -18,6 +18,7 @@
 from .transverse_generators import generate_2D_pencil
 from .transverse_generators import generate_2D_pencil_with_absolute_cut
 from .transverse_generators import generate_2D_gaussian
+from .transverse_generators import hypersphere_2D, hypersphere_4D, hypersphere_6D
 
 from .longitudinal import generate_longitudinal_coordinates
 from .longitudinal.generate_longitudinal import _characterize_line

xpart/build_particles.py

@@ -53,7 +53,7 @@ def build_particles(_context=None, _buffer=None, _offset=None, _capacity=None,
                       R_matrix=None,
                       W_matrix=None,
                       method=None,
-                      nemitt_x=None, nemitt_y=None,
+                      nemitt_x=None, nemitt_y=None,nemitt_zeta = None,
                       scale_with_transverse_norm_emitt=None,
                       weight=None,
                       **kwargs, # They are passed to the twiss
@@ -258,7 +258,11 @@ def build_particles(_context=None, _buffer=None, _offset=None, _capacity=None,
             gemitt_y = (nemitt_y / particle_ref._xobject.beta0[0]
                         / particle_ref._xobject.gamma0[0])
 
-        gemitt_zeta = 1
+        if nemitt_zeta is None:
+            gemitt_zeta = 1
+        else:
+            gemitt_zeta = (nemitt_zeta / particle_ref._xobject.beta0[0]
+                        / particle_ref._xobject.gamma0[0])
 
 
         n_constraints = sum([vv is not None for vv in [x, x_norm, px, px_norm,

xpart/transverse_generators/__init__.py

@@ -7,3 +7,4 @@
 from .polar import generate_2D_uniform_circular_sector
 from .pencil import generate_2D_pencil, generate_2D_pencil_with_absolute_cut
 from .gaussian import generate_2D_gaussian
+from .hypersphere import hypersphere_2D, hypersphere_4D, hypersphere_6D

xpart/transverse_generators/hypersphere.py

@@ -0,0 +1,135 @@
+
+import numpy as np
+
+
+
+
+def hypersphere(N, D, r=1, rng_seed = 0, surface=False ,unpack = False):
+    '''
+    Generate points uniformly distributed inside or on the surface of an N-dimensional hypersphere.
+    Adapted from : https://baezortega.github.io/2018/10/14/hypersphere-sampling/
+
+    Parameters
+    ----------
+    N : int
+        Number of particles to generate.
+    D : int
+        Dimension of the hypersphere.
+    r : float or list
+        Radius of the hypersphere. If a list, specifies radii for anisotropic scaling.
+    rng_seed : int
+        Seed for the random number generator for reproducibility.
+    surface : bool
+        If True, points will be generated on the surface. If False, points will be generated inside the hypersphere.
+    unpack : bool
+        If True, returns individual arrays for each dimension. If False, returns a single array with shape (N, D).
+
+    Returns
+    -------
+    samples : np.ndarray or tuple of np.ndarray
+        Generated points. Shape is (N, D) if unpack is False, otherwise D arrays of shape (N,).
+    '''
+    # Set the random seed for reproducibility
+    rng = np.random.default_rng(int(rng_seed))
+
+    # Sample D vectors of N Gaussian coordinates
+    N = int(N)
+    D = int(D)
+    samples = rng.standard_normal(size = (N, D))
+
+    # Normalise all distances (radii) to 1
+    radii = np.sqrt(np.sum(samples ** 2, axis=1))[:,np.newaxis]
+    samples = samples / radii
+
+    # Sample N radii with exponential distribution (unless points are to be on the surface)
+    if not surface:
+        new_radii = np.random.uniform(low=0.0, high=1.0, size=(N, 1)) ** (1 / D)
+        samples = samples * new_radii
+
+    # Scale the samples to the desired radius
+    if isinstance(r,list):
+        r = np.array(r)[np.newaxis,:]
+    elif isinstance(r,type(np.array([]))):
+        assert False, 'r should be float or list'
+    samples = samples * r
+
+    if not unpack:
+        return samples
+    else:
+        return samples.T
+
+
+
+def hypersphere_2D(num_particles,r = 1, rng_seed = 0):
+    '''
+    Generate points uniformly distributed inside a 2-dimensional hypersphere (circle).
+
+    Parameters
+    ----------
+    num_particles : int
+        Number of particles to generate.
+    r : float
+        Radius of the circle.
+    rng_seed : int
+        Seed for the random number generator for reproducibility.
+
+    Returns
+    -------
+    x_norm : np.ndarray
+        x-coordinates of the generated points.
+    px_norm : np.ndarray
+        y-coordinates of the generated points.
+    '''
+    x_norm , px_norm  = hypersphere(num_particles,D=2,r=r, rng_seed=rng_seed, surface = False,unpack=True)
+
+    return x_norm, px_norm
+
+
+def hypersphere_4D(num_particles,rx =1,ry =1, rng_seed = 0):
+    '''
+    Generate points uniformly distributed inside a 4-dimensional hypersphere with anisotropic scaling.
+
+    Parameters
+    ----------
+    num_particles : int
+        Number of particles to generate.
+    rx : float
+        Scaling factor for the x and px dimensions.
+    ry : float
+        Scaling factor for the y and py dimensions.
+    rng_seed : int
+        Seed for the random number generator for reproducibility.
+
+    Returns
+    -------
+    x_norm, px_norm, y_norm, py_norm : np.ndarray
+        Coordinates of the generated points in the 4-dimensional space.
+    '''
+
+    x_norm , px_norm , y_norm, py_norm = hypersphere(num_particles,D=4,r=[rx,rx,ry,ry], rng_seed=rng_seed, surface = False,unpack=True)
+
+    return x_norm , px_norm , y_norm, py_norm
+
+
+def hypersphere_6D(num_particles,rx =1,ry =1, rzeta=1, rng_seed = 0):
+    '''
+    Generate points uniformly distributed inside a 6-dimensional hypersphere with anisotropic scaling.
+
+    Parameters
+    ----------
+    num_particles : int
+        Number of particles to generate.
+    rx, ry, rzeta : float
+        Scaling factors for the x, px, y, py, zeta, and pzeta dimensions respectively.
+    rng_seed : int
+        Seed for the random number generator for reproducibility.
+
+    Returns
+    -------
+    x_norm, px_norm, y_norm, py_norm, zeta_norm, pzeta_norm : np.ndarray
+        Coordinates of the generated points in the 6-dimensional space.
+    '''
+
+    x_norm , px_norm , y_norm, py_norm, zeta_norm, pzeta_norm = hypersphere(num_particles,D=6,r=[rx,rx,ry,ry,rzeta,rzeta], rng_seed=rng_seed, surface = False,unpack=True)
+
+    return x_norm , px_norm , y_norm, py_norm, zeta_norm, pzeta_norm