From 31641ec587ac0d183a0cf8ef23dced44393577f6 Mon Sep 17 00:00:00 2001 From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com> Date: Wed, 15 Nov 2023 13:30:32 -0800 Subject: [PATCH] add equal_weights option to resample --- pmd_beamphysics/interfaces/genesis.py | 5 +++-- pmd_beamphysics/particles.py | 4 ++-- pmd_beamphysics/statistics.py | 31 ++++++++++++++++----------- 3 files changed, 23 insertions(+), 17 deletions(-) diff --git a/pmd_beamphysics/interfaces/genesis.py b/pmd_beamphysics/interfaces/genesis.py index d566fdf..447d1c8 100644 --- a/pmd_beamphysics/interfaces/genesis.py +++ b/pmd_beamphysics/interfaces/genesis.py @@ -398,7 +398,8 @@ def write_genesis4_distribution(particle_group, If particles are at different z, they will be drifted to the same z, because the output should have different times. - If any of the weights are different, the bunch will be resampled. + If any of the weights are different, the bunch will be resampled + to have equal weights. Note that this can be very slow for a large number of particles. """ @@ -423,7 +424,7 @@ def write_genesis4_distribution(particle_group, n = len(P) if verbose: print(f'Resampling {n} weighted particles') - P = P.resample(n) + P = P.resample(n, equal_weights=True) for k in ['x', 'xp', 'y', 'yp', 't']: h5[k] = P[k] diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py index e036a5f..f0e3d66 100644 --- a/pmd_beamphysics/particles.py +++ b/pmd_beamphysics/particles.py @@ -1064,8 +1064,8 @@ def copy(self): return deepcopy(self) @functools.wraps(resample_particles) - def resample(self, n=0): - data = resample_particles(self, n) + def resample(self, n=0, equal_weights=False): + data = resample_particles(self, n, equal_weights=equal_weights) return ParticleGroup(data=data) # Internal sorting diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index fb8bc6f..0be5504 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -460,7 +460,7 @@ def slice_statistics(particle_group, keys=['mean_z'], n_slice=40, slice_key=Non -def resample_particles(particle_group, n=0): +def resample_particles(particle_group, n=0, equal_weights=False): """ Resamples a ParticleGroup randomly. @@ -479,6 +479,9 @@ def resample_particles(particle_group, n=0): Number to resample. If n = 0, this will use all particles. + equal_weights: bool, default = False + If True, will ensure that all particles have equal weights. + Returns ------- data: dict of ParticleGroup data @@ -498,19 +501,21 @@ def resample_particles(particle_group, n=0): ixlist = np.random.choice(n_old, n, replace=False) weight = np.full(n, particle_group.charge/n) - # variable weights + # variable weights found + elif equal_weights or n != n_old: + # From SciPy example: + # https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.rv_discrete.html#scipy.stats.rv_discrete + pk = weight / np.sum(weight) # Probabilities + xk = np.arange(len(pk)) # index + ixsampler = scipy_stats.rv_discrete(name='ixsampler', values=(xk, pk)) + ixlist = ixsampler.rvs(size=n) + weight = np.full(n, particle_group.charge/n) + else: - if n == n_old: - ixlist = np.random.choice(n_old, n, replace=False) - weight = weight[ixlist] #just scramble - else: - # From SciPy example: - # https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.rv_discrete.html#scipy.stats.rv_discrete - pk = weight / np.sum(weight) # Probabilities - xk = np.arange(len(pk)) # index - ixsampler = scipy_stats.rv_discrete(name='ixsampler', values=(xk, pk)) - ixlist = ixsampler.rvs(size=n) - weight = np.full(n, particle_group.charge/n) + assert n == n_old + ixlist = np.random.choice(n_old, n, replace=False) + weight = weight[ixlist] #just scramble + data = {} for key in particle_group._settable_array_keys: