From 52fceec2b7bf3c48f946d175ead8d7e36b3eb546 Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Wed, 16 Aug 2023 11:57:36 -0500 Subject: [PATCH 1/8] add normalized longitudinal coordinates --- pmd_beamphysics/particles.py | 19 +++++++++++++++++-- pmd_beamphysics/statistics.py | 26 ++++++++++++++++++-------- 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py index 559491f..b217104 100644 --- a/pmd_beamphysics/particles.py +++ b/pmd_beamphysics/particles.py @@ -545,8 +545,23 @@ def py_bar(self): @property def Jy(self): """Normalized amplitude J in the y-py plane""" - return particle_amplitude(self, 'y') - + return particle_amplitude(self, 'y') + + @property + def z_bar(self): + """Normalized y in units of sqrt(m)""" + return normalized_particle_coordinate(self, 'z') + + @property + def pz_bar(self): + """Normalized py in units of sqrt(m)""" + return normalized_particle_coordinate(self, 'pz') + + @property + def Jz(self): + """Normalized amplitude J in the y-py plane""" + return (self.z_bar**2 + self.pz_bar**2) / 2 + def delta(self, key): """Attribute (array) relative to its mean""" return self[key] - self.avg(key) diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index 679270b..a3a3c9e 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -1,5 +1,7 @@ import numpy as np from scipy import stats as scipy_stats +from scipy.constants import c + def norm_emit_calc(particle_group, planes=['x']): """ @@ -344,7 +346,7 @@ def normalized_particle_coordinate(particle_group, key, twiss=None, mass_normali and the Intended use is for key to be one of: - x, px, y py + x, px, y, py, z, pz and the corresponding normalized coordinates are named with suffix _bar, i.e.: x_bar, px_bar, y_bar, py_bar @@ -378,16 +380,24 @@ def normalized_particle_coordinate(particle_group, key, twiss=None, mass_normali momentum = False key1 = key key2 = 'p'+key - - x = particle_group[key1] - - if mass_normalize: - # Note: do not do /=, because this will replace the ParticleGroup's internal array! - p = particle_group[key2]/particle_group.mass + + # get position and momentum + # handle longitudinal case + if key1 == "z": + x = particle_group.delta("t")*c + else: + x = particle_group[key1] + + if key2 == "pz": + p = particle_group.delta("pz") else: p = particle_group[key2] - + if mass_normalize: + # Note: do not do /=, because this will replace the ParticleGroup's internal array! + p = p / particle_group.mass + + # User could supply twiss if not twiss: sigma_mat2 = np.cov(x, p, aweights=particle_group.weight) From b187aa02f3eaf3cd72c9c96a210ede4fa3efd242 Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 18 Aug 2023 10:26:22 -0500 Subject: [PATCH 2/8] Revert "add normalized longitudinal coordinates" This reverts commit 52fceec2b7bf3c48f946d175ead8d7e36b3eb546. --- pmd_beamphysics/particles.py | 19 ++----------------- pmd_beamphysics/statistics.py | 26 ++++++++------------------ 2 files changed, 10 insertions(+), 35 deletions(-) diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py index b217104..559491f 100644 --- a/pmd_beamphysics/particles.py +++ b/pmd_beamphysics/particles.py @@ -545,23 +545,8 @@ def py_bar(self): @property def Jy(self): """Normalized amplitude J in the y-py plane""" - return particle_amplitude(self, 'y') - - @property - def z_bar(self): - """Normalized y in units of sqrt(m)""" - return normalized_particle_coordinate(self, 'z') - - @property - def pz_bar(self): - """Normalized py in units of sqrt(m)""" - return normalized_particle_coordinate(self, 'pz') - - @property - def Jz(self): - """Normalized amplitude J in the y-py plane""" - return (self.z_bar**2 + self.pz_bar**2) / 2 - + return particle_amplitude(self, 'y') + def delta(self, key): """Attribute (array) relative to its mean""" return self[key] - self.avg(key) diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index a3a3c9e..679270b 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -1,7 +1,5 @@ import numpy as np from scipy import stats as scipy_stats -from scipy.constants import c - def norm_emit_calc(particle_group, planes=['x']): """ @@ -346,7 +344,7 @@ def normalized_particle_coordinate(particle_group, key, twiss=None, mass_normali and the Intended use is for key to be one of: - x, px, y, py, z, pz + x, px, y py and the corresponding normalized coordinates are named with suffix _bar, i.e.: x_bar, px_bar, y_bar, py_bar @@ -380,24 +378,16 @@ def normalized_particle_coordinate(particle_group, key, twiss=None, mass_normali momentum = False key1 = key key2 = 'p'+key - - # get position and momentum - # handle longitudinal case - if key1 == "z": - x = particle_group.delta("t")*c - else: - x = particle_group[key1] - - if key2 == "pz": - p = particle_group.delta("pz") - else: - p = particle_group[key2] + + x = particle_group[key1] if mass_normalize: # Note: do not do /=, because this will replace the ParticleGroup's internal array! - p = p / particle_group.mass - - + p = particle_group[key2]/particle_group.mass + else: + p = particle_group[key2] + + # User could supply twiss if not twiss: sigma_mat2 = np.cov(x, p, aweights=particle_group.weight) From a0ad2fc3216de2c378b49cb7e108891c13c6e8b9 Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 18 Aug 2023 10:41:59 -0500 Subject: [PATCH 3/8] add ability to calculate 6D particle amplitudes --- pmd_beamphysics/statistics.py | 60 ++++++++++++++++++++++++----------- 1 file changed, 42 insertions(+), 18 deletions(-) diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index 679270b..461c62e 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -308,32 +308,56 @@ def particle_amplitude(particle_group, plane='x', twiss=None, mass_normalize=Tru Plane should be: 'x' for the x, px plane - 'y' for the y, py plane + 'y' for the y, py plane + `all` for the 6D phase space Other planes will work, but please check that the units make sense. - + If mass_normalize (default=True), the momentum will be divided by the mass, so that the units are sqrt(m). See: normalized_particle_coordinate """ - x = particle_group[plane] - key2 = 'p'+plane - - if mass_normalize: - # Note: do not do /=, because this will replace the ParticleGroup's internal array! - p = particle_group[key2]/particle_group.mass + if plane != "all": + x = particle_group[plane] + key2 = 'p'+plane + + if mass_normalize: + # Note: do not do /=, because this will replace the ParticleGroup's internal array! + p = particle_group[key2]/particle_group.mass + else: + p = particle_group[key2] + + # User could supply twiss + if not twiss: + sigma_mat2 = np.cov(x, p, aweights=particle_group.weight) + twiss = twiss_calc(sigma_mat2) + + J = amplitude_calc(x, p, beta=twiss['beta'], alpha=twiss['alpha']) else: - p = particle_group[key2] - - # User could supply twiss - if not twiss: - sigma_mat2 = np.cov(x, p, aweights=particle_group.weight) - twiss = twiss_calc(sigma_mat2) - - J = amplitude_calc(x, p, beta=twiss['beta'], alpha=twiss['alpha']) + vnames = ["x", "px", "y", "py", "t", "pz"] + data = np.copy(np.stack([particle_group[name] for name in vnames]).T) + + # get delta t and pz + data[:, -2] = data[:, -2] - np.mean(data[:, -2]) + data[:, -1] = data[:, -1] - np.mean(data[:, -1]) + + if mass_normalize: + for i in [1, 3, 5]: + data[:, i] = data[:, i] / particle_group.mass + + if twiss is not None: + raise ValueError("cannot specify twiss parameters for 6D amplitude " + "calculation") + + cov = np.cov(data.T, aweights=particle_group.weight) + + # transform particle coordinates into normalized coordinates using inverse of + # Cholesky decomp + t_data = (np.linalg.inv(np.linalg.cholesky(cov)) @ data.T).T + + J = np.linalg.norm(t_data, axis=1) return J - - + def normalized_particle_coordinate(particle_group, key, twiss=None, mass_normalize=True): """ From 811240d06dceacaf449bbc70a83bb90b7f5e352c Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 18 Aug 2023 10:46:22 -0500 Subject: [PATCH 4/8] modularize functions --- pmd_beamphysics/statistics.py | 36 +++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index 461c62e..35f98eb 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -333,31 +333,31 @@ def particle_amplitude(particle_group, plane='x', twiss=None, mass_normalize=Tru J = amplitude_calc(x, p, beta=twiss['beta'], alpha=twiss['alpha']) else: - vnames = ["x", "px", "y", "py", "t", "pz"] - data = np.copy(np.stack([particle_group[name] for name in vnames]).T) + t_data = normalized_6D_coordinates(particle_group, mass_normalize) - # get delta t and pz - data[:, -2] = data[:, -2] - np.mean(data[:, -2]) - data[:, -1] = data[:, -1] - np.mean(data[:, -1]) + J = np.linalg.norm(t_data, axis=1) + + return J - if mass_normalize: - for i in [1, 3, 5]: - data[:, i] = data[:, i] / particle_group.mass +def normalized_6D_coordinates(particle_group, mass_normalize=True): + vnames = ["x", "px", "y", "py", "t", "pz"] + data = np.copy(np.stack([particle_group[name] for name in vnames]).T) - if twiss is not None: - raise ValueError("cannot specify twiss parameters for 6D amplitude " - "calculation") + # get delta t and pz + data[:, -2] = data[:, -2] - np.mean(data[:, -2]) + data[:, -1] = data[:, -1] - np.mean(data[:, -1]) - cov = np.cov(data.T, aweights=particle_group.weight) + if mass_normalize: + for i in [1, 3, 5]: + data[:, i] = data[:, i] / particle_group.mass - # transform particle coordinates into normalized coordinates using inverse of - # Cholesky decomp - t_data = (np.linalg.inv(np.linalg.cholesky(cov)) @ data.T).T + cov = np.cov(data.T, aweights=particle_group.weight) - J = np.linalg.norm(t_data, axis=1) - - return J + # transform particle coordinates into normalized coordinates using inverse of + # Cholesky decomp + t_data = (np.linalg.inv(np.linalg.cholesky(cov)) @ data.T).T + return t_data def normalized_particle_coordinate(particle_group, key, twiss=None, mass_normalize=True): """ From c159478bb9df7937e7c3a786cfd965e1e1153a94 Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 18 Aug 2023 11:24:19 -0500 Subject: [PATCH 5/8] add docstring --- pmd_beamphysics/statistics.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index 35f98eb..0ef4b3c 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -340,6 +340,15 @@ def particle_amplitude(particle_group, plane='x', twiss=None, mass_normalize=Tru return J def normalized_6D_coordinates(particle_group, mass_normalize=True): + """ + Compute coordinates in 6D phase space normalized to the measured beam + matrix. Default behavior is to normalize momenta by mass such that the normalized + unit is in 1/sqrt(m). + + :param particle_group: + :param mass_normalize: + :return: + """ vnames = ["x", "px", "y", "py", "t", "pz"] data = np.copy(np.stack([particle_group[name] for name in vnames]).T) From 2e7bb9fd3637112761b1f6b09fc329c117bcb107 Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 18 Aug 2023 11:28:25 -0500 Subject: [PATCH 6/8] handle t/z coordinate definition, add equation to docstring --- pmd_beamphysics/statistics.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index 0ef4b3c..f28aa1b 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -345,11 +345,16 @@ def normalized_6D_coordinates(particle_group, mass_normalize=True): matrix. Default behavior is to normalize momenta by mass such that the normalized unit is in 1/sqrt(m). + If x is a vector of particle coordinates in 6D space: + x_bar = S^{-1}x where SS^T = cov(x, x) + :param particle_group: :param mass_normalize: :return: """ - vnames = ["x", "px", "y", "py", "t", "pz"] + longitudinal_coordinate = "t" if particle_group.in_z_coordinates else "z" + + vnames = ["x", "px", "y", "py", longitudinal_coordinate, "pz"] data = np.copy(np.stack([particle_group[name] for name in vnames]).T) # get delta t and pz From e1d1b8a296153986aa18686cd6b6094ca57f42d0 Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 18 Aug 2023 11:53:12 -0500 Subject: [PATCH 7/8] Update statistics.py --- pmd_beamphysics/statistics.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index f28aa1b..507d1a0 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -346,14 +346,24 @@ def normalized_6D_coordinates(particle_group, mass_normalize=True): unit is in 1/sqrt(m). If x is a vector of particle coordinates in 6D space: - x_bar = S^{-1}x where SS^T = cov(x, x) + x_bar = S^{-1}x where SS^T = cov(x, x) and x_bar is the coordinates in normalized + space. + + Parameters + ---------- + particle_group: ParticleGroup + Particle group to be transformed + mass_normalize: bool, defualt: True + Flag to divide by particle mass to return nomalized momentum units in 1/sqrt(m). + + Returns + ------- + res : ndarray + Particle coordinates in normalized phase space - :param particle_group: - :param mass_normalize: - :return: """ longitudinal_coordinate = "t" if particle_group.in_z_coordinates else "z" - + vnames = ["x", "px", "y", "py", longitudinal_coordinate, "pz"] data = np.copy(np.stack([particle_group[name] for name in vnames]).T) From 8545703af1e886be8990ed801c2d9b243ebbd70a Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Thu, 24 Oct 2024 09:34:53 -0500 Subject: [PATCH 8/8] Update pmd_beamphysics/statistics.py Co-authored-by: Christopher M. Pierce --- pmd_beamphysics/statistics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pmd_beamphysics/statistics.py b/pmd_beamphysics/statistics.py index 507d1a0..f2035d5 100644 --- a/pmd_beamphysics/statistics.py +++ b/pmd_beamphysics/statistics.py @@ -379,7 +379,7 @@ def normalized_6D_coordinates(particle_group, mass_normalize=True): # transform particle coordinates into normalized coordinates using inverse of # Cholesky decomp - t_data = (np.linalg.inv(np.linalg.cholesky(cov)) @ data.T).T + t_data = np.linalg.solve(np.linalg.cholesky(cov), data.T).T return t_data