6D Amplitude Calculation

pmd_beamphysics/statistics.py

@@ -323,31 +323,82 @@ 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
+        `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 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]
+        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)
+        # 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:
+        t_data = normalized_6D_coordinates(particle_group, mass_normalize)
 
-    J = amplitude_calc(x, p, beta=twiss["beta"], alpha=twiss["alpha"])
+        J = np.linalg.norm(t_data, axis=1)
 
     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).
+
+    If x is a vector of particle coordinates in 6D space:
+    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
+
+    """
+    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
+    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
+
+    cov = np.cov(data.T, aweights=particle_group.weight)
+
+    # transform particle coordinates into normalized coordinates using inverse of
+    # Cholesky decomp
+    t_data = np.linalg.solve(np.linalg.cholesky(cov), data.T).T
+
+    return t_data
+
+
 def normalized_particle_coordinate(
     particle_group, key, twiss=None, mass_normalize=True
 ):