# define the same transfer matrix as we calculate the twiss, ignore the dispersion term

import numpy as np

def D(l):
    return np.array([[1,l,0,0],
                    [0,1,0,0],
                    [0,0,1,l],
                    [0,0,0,1]])

def QF(l,k):
    sk=np.sqrt(k)
    skl=sk*l
    return np.array([[np.cos(skl),np.sin(skl)/sk,0,0],
                     [-sk*np.sin(skl),np.cos(skl),0,0],
                     [0,0,np.cosh(skl),np.sinh(skl)/sk],
                     [0,0,sk*np.sinh(skl),np.cosh(skl)]])

def QD(l,k):
    sk=np.sqrt(np.abs(k))
    skl=sk*l
    return np.array([[np.cosh(skl),np.sinh(skl)/sk,0,0],
                     [sk*np.sinh(skl),np.cosh(skl),0,0],
                     [0,0,np.cos(skl),np.sin(skl)/sk],
                     [0,0,-sk*np.sin(skl),np.cos(skl)]])

# define a function return matrix for sector magnet(recommand rho and angle form):
def SBend(r, a):
    return np.array([[np.cos(a),    r*np.sin(a),   0,0],
                     [-np.sin(a)/r, np.cos(a),     0,0],
                     [0,0,1,r*a], # vertical as dirft space
                     [0,0,0,1]])

# return element matrix by its type and parameters
def element_matrix(*args):
    if args[0] == 'D':
        l=args[1]
        return D(l)
    elif args[0] == 'Q':
        l=args[1]
        k=args[2]
        if k>0:
            return QF(l,k)
        elif k<0:
            return QD(l,k)
        else:
            return [D(l),D(l),l] # if k=0, return as drift
    elif args[0] == 'SBend':
        l=args[1]
        angle=args[2]
        rho=l/angle
        return SBend(rho,angle)
    else:
        print('Unknown element type!')
        return None

# define a function doing the do production for a list of given matrix:
def Mdot(ml):
    mt = ml[0] # temp matrix to store the results
    if len(ml) == 1: # return if there is only one matrix
        return mt
    else:
        # do the dot production one by one from the 2nd element
        for i in range(1,len(ml)):
            mt = np.dot(ml[i],mt)
    return mt

# the rcs r1 lattice definition
rcs_r1=[["D",5.5,0],
["Q",0.41,0.7406566],
["D",0.8,0],
["Q",0.9,-0.579364096],
["D",1.15,0],
["Q",0.41,0.664196402],
["D",3.8,0],
["SBend",2.1,0.261799388],
["D",1.2,0],
["SBend",2.1,0.261799388],
["D",1.3,0],
["Q",0.45,0.570558026],
["D",1.3,0],
["Q",0.9,-0.579364096],
["D",0.8,0],
["Q",0.62,0.608457761],
["D",0.9,0],
["SBend",2.1,0.261799388],
["D",3.5,0],
["SBend",2.1,0.261799388],
["D",0.9,0],
["Q",0.62,0.608457761],
["D",0.8,0],
["Q",0.9,-0.579364096],
["D",1.3,0],
["Q",0.45,0.570558026],
["D",1.3,0],
["SBend",2.1,0.261799388],
["D",1.2,0],
["SBend",2.1,0.261799388],
["D",3.8,0],
["Q",0.41,0.664196402],
["D",1.15,0],
["Q",0.9,-0.579364096],
["D",0.8,0],
["Q",0.41,0.7406566],
["D",5.5,0]]

# A quick test to compared with our twiss script about the matrix element
ml=[]
for i,e in enumerate(rcs_r1):
    ml=[element_matrix(e[0],e[1],e[2]) for e in rcs_r1]

print("Transfer matrix for r1:\n",Mdot(ml))
print("Transfer matrix for rcs(4*r1):\n",Mdot(ml*4))

Transfer matrix for r1:
 [[-0.43160468  6.65011376  0.          0.        ]
 [-0.12236143 -0.43160468  0.          0.        ]
 [ 0.          0.          0.8359187   2.53133998]
 [ 0.          0.         -0.11900413  0.8359187 ]]
Transfer matrix for rcs(4*r1):
 [[-0.21265114  7.2035042   0.          0.        ]
 [-0.13254375 -0.21265114  0.          0.        ]
 [ 0.          0.         -0.68395546  3.36460171]
 [ 0.          0.         -0.15817769 -0.68395546]]

# a simple track function for one particle, p is in the shape of np.array([[x],[xp],[y],[yp]])
def track_single(tm,p):
    return np.dot(tm,p)
# track a bunch of particles
def track(tm,p):
    return np.array([track_single(tm,p[i]) for i in range(len(p))])

# test particle at position x=1 and y=1, both xp and yp are zero
x0=np.array([1,0,1,0]).reshape(4,1)
# track for 100 times
r=[]
TM=Mdot(ml*4)
r.append(x0) # store the initial distribution first
for i in range(100):
    x0=track_single(TM,x0)
    r.append(x0)

print("Single particle trace should paint the envelop defined by the local twiss paramters")

from matplotlib import pyplot as plt

# Extract the pairs (0, 1) and (2, 3) from the arrays
pair_01 = np.array([[r[i][0][0], r[i][1][0]] for i in range(len(r))])
pair_23 = np.array([[r[i][2][0], r[i][3][0]] for i in range(len(r))])

# plot trace on x-xp and y-yp with differen colors
plt.scatter(pair_01[:, 0], pair_01[:, 1], color='r', label='x-xp')
plt.scatter(pair_23[:, 0], pair_23[:, 1], color='b', label='y-yp')

plt.xlabel('x or y')
plt.ylabel('xp or yp')
plt.legend()
plt.show()

Single particle trace should paint the envelop defined by the local twiss paramters

n=1000
# generate a bunch of particles in the form of [x,xp,y,yp] uniformly distributed within -1 and 1
p=np.random.rand(n,4)
p=p*2-1
p=p.reshape(n,4,1)

# Plot x-xp, y-yp, and x-y on three subplots in a single figure
fig, axs = plt.subplots(1, 3, figsize=(18, 6))

# Plot x-xp
axs[0].scatter(p[:, 0], p[:, 1], s=5, color='b', alpha=0.5)
axs[0].set_xlabel('x')
axs[0].set_ylabel('xp')
axs[0].set_title('x versus xp')
axs[0].grid(True)

# Plot y-yp
axs[1].scatter(p[:, 2], p[:, 3], s=5, color='g', alpha=0.5)
axs[1].set_xlabel('y')
axs[1].set_ylabel('yp')
axs[1].set_title('y versus yp')
axs[1].grid(True)

# Plot x-y
axs[2].scatter(p[:, 0], p[:, 2], s=5, color='r', alpha=0.5)
axs[2].set_xlabel('x')
axs[2].set_ylabel('y')
axs[2].set_title('x versus y')
axs[2].grid(True)

plt.tight_layout()
plt.show()

# track the whole bunch for 100 turns and store their distribution
r=[]
r.append(p) # store initial distribution
for i in range(1,1000):
    p=track(TM,p)
    r.append(p)

# calculate the emittance for a given distribution: e=sqrt(mean(x^2)*mean(xp^2)-mean(x*xp))
def distribution_parameters(p):
    ex=np.sqrt(np.mean((p[:,0]-np.mean(p[:,0]))**2)*np.mean((p[:,1]-np.mean(p[:,1]))**2)-np.mean(p[:,0]*p[:,1])**2)
    ey=np.sqrt(np.mean((p[:,2]-np.mean(p[:,2]))**2)*np.mean((p[:,3]-np.mean(p[:,3]))**2)-np.mean(p[:,2]*p[:,3])**2)
    # beam size given as standard deviation
    sx=np.std(p[:,0])
    sy=np.std(p[:,2])
    return ex/np.pi,ey/np.pi,sx,sy

# a figure for distribution plot
def distribution_plot(p,xlim=[-10,10],ylim=[-10,10],scale=1000):
    # Plot x-xp, y-yp, and x-y on three subplots in a single figure
    fig, axs = plt.subplots(1, 3, figsize=(18, 6))

    ex,ey,sx,sy=distribution_parameters(p)

    # Plot x-xp
    axs[0].set_xlim(xlim)
    axs[0].set_ylim(ylim)
    axs[0].scatter(p[:, 0]*scale, p[:, 1]*scale, s=5, color='b', alpha=0.5)
    axs[0].set_xlabel('x')
    axs[0].set_ylabel('xp')
    axs[0].set_title('x-xp: {:.3f} pi.mm.mrad'.format(ex*scale*scale))
    axs[0].grid(True)

    # Plot y-yp
    axs[1].set_xlim(xlim)
    axs[1].set_ylim(ylim)
    axs[1].scatter(p[:, 2]*scale, p[:, 3]*scale, s=5, color='g', alpha=0.5)
    axs[1].set_xlabel('y')
    axs[1].set_ylabel('yp')
    axs[1].set_title('y-yp: {:.3f} pi.mm.mrad'.format(ey*scale*scale))
    axs[1].grid(True)

    # Plot x-y
    axs[1].set_xlim(xlim)
    axs[1].set_ylim(ylim)
    axs[2].scatter(p[:, 0]*scale, p[:, 2]*scale, s=5, color='r', alpha=0.5)
    axs[2].set_xlabel('x')
    axs[2].set_ylabel('y')
    axs[2].set_title('x versus y')
    axs[2].grid(True)

    plt.tight_layout()
    plt.show()

distribution_plot(r[99])

n=10000
# generate a bunch of particles in the form of [x,xp,y,yp] with gaussian distribution and sigma =1
p=np.random.randn(n,4)
# p=p.reshape(n,4,1)

# distribution_plot(p,[-10,10],[-10,10])

# track the whole bunch for 100 turns and store their distribution
rn=[]
pn=[]
rn.append(p) # store initial distribution
ex0,ey0,sx0,sy0=distribution_parameters(p)
pn.append([0,ex0,ey0,sx0,sy0])
for i in range(1,1000):
    p=track(TM,p)
    rn.append(p)
    ex,ey,sx,sy=distribution_parameters(p)
    pn.append([i+1,ex,ey,sx,sy])

# plot ex,ey,sx and sy in terms of i for pn

# Extract the data for plotting
turns = [entry[0] for entry in pn]
parameters = {
    'Horizontal Emittance (ex)': [entry[1] for entry in pn],
    'Vertical Emittance (ey)': [entry[2] for entry in pn],
    'Horizontal Beam Size (sx)': [entry[3] for entry in pn],
    'Vertical Beam Size (sy)': [entry[4] for entry in pn]
}

# Create a figure and axis
fig, ax = plt.subplots(figsize=(12, 8))

# Plot each parameter
for param_name, param_values in parameters.items():
    linestyle = '--' if param_name in ['Horizontal Emittance (ex)', 'Vertical Emittance (ey)'] else '-'
    ax.plot(turns, param_values, label=param_name, linestyle=linestyle)

# Set labels and title
ax.set_xlabel('Turn Number (i)')
ax.set_ylabel('Parameter Value')
ax.set_title('Evolution of Emittance and Beam Size')

# Display legend and grid
ax.legend()
ax.grid(True)

# Show the plot
plt.show()

import matplotlib.animation as animation

fig, axs = plt.subplots(1, 3, figsize=(15, 5))

frames=rn
n_frames=len(frames)
xlim=[-30,30]
ylim=[-30,30]
# Initialization function: plot the first frame
def init():
    axs[0].set_xlim(xlim)
    axs[0].set_ylim(ylim)
    axs[0].scatter(frames[0][:, 0], frames[0][:, 1], s=5, color='b', alpha=0.5)
    axs[0].set_xlabel('x')
    axs[0].set_ylabel('xp')
    axs[0].set_title('x versus xp')
    
    axs[1].set_xlim(xlim)
    axs[1].set_ylim(ylim)
    axs[1].scatter(frames[0][:, 2], frames[0][:, 3], s=5, color='g', alpha=0.5)
    axs[1].set_xlabel('y')
    axs[1].set_ylabel('yp')
    axs[1].set_title('y versus yp')

    axs[2].set_xlim(xlim)
    axs[2].set_ylim(ylim)
    axs[2].scatter(frames[0][:, 0], frames[0][:, 2], s=5, color='r', alpha=0.5)
    axs[2].set_xlabel('x')
    axs[2].set_ylabel('y')
    axs[2].set_title('x versus y')

# Update function: update the scatter plot for each frame
def update(frame):

    ex,ey,sx,sy=distribution_parameters(frames[frame])

    axs[0].clear()
    axs[0].set_xlim(xlim)
    axs[0].set_ylim(ylim)
    axs[0].scatter(frames[frame][:, 0], frames[frame][:, 1], s=5, color='b', alpha=0.5)
    axs[0].set_xlabel('x')
    axs[0].set_ylabel('xp')
    axs[0].set_title('x-xp: {:.3f} pi.mm.mrad'.format(ex))

    axs[1].clear()
    axs[1].set_xlim(xlim)
    axs[1].set_ylim(ylim)
    axs[1].scatter(frames[frame][:, 2], frames[frame][:, 3], s=5, color='g', alpha=0.5)
    axs[1].set_xlabel('y')
    axs[1].set_ylabel('yp')
    axs[1].set_title('y-yp: {:.3f} pi.mm.mrad'.format(ey))

    axs[2].clear()
    axs[2].set_xlim(xlim)
    axs[2].set_ylim(ylim)
    axs[2].scatter(frames[frame][:, 0], frames[frame][:, 2], s=5, color='r', alpha=0.5)
    axs[2].set_xlabel('x')
    axs[2].set_ylabel('y')
    axs[2].set_title('x - y')

    fig.suptitle('TURN {}'.format(frame), fontsize=16)  # Set the overall title for the animation

ani = animation.FuncAnimation(fig, update, frames=n_frames, init_func=init, blit=False)

ani.save('particle_distribution_n.gif', writer='pillow', fps=10)
plt.show()

distribution_plot(rn[0],[-30,30],[-30,30])

# matched beam generation
print("Notice that at the staring point, alpha_x,y is zero, \n which means that the distribution is a typical ellipse.ex=sqrt(sx^2*sxp^2-sxxp^2)")
ex0=ey0=1.0e-6 # pi.mm.mrad
bx0=7.372117913456016
by0=4.612052201102539
sx0=np.sqrt(bx0*ex0*np.pi)
sxp0=np.sqrt(ex0*np.pi/bx0)
# sxp0=ex0/sx0
sy0=np.sqrt(by0*ey0*np.pi)
syp0=np.sqrt(ey0*np.pi/by0)
# syp0=ey0/sy0

print("sx0,sxp0,sy0,syp0",[sx0*1000,sxp0*1000,sy0*1000,syp0*1000])

num_particles =10000
# Covariance matrix
cov_x = [[sx0**2, 0],
              [0, sxp0**2]]
cov_y = [[sy0**2, 0],
              [0, syp0**2]]

# Generate random samples from a multivariate normal distribution
mean = [0, 0]
x, xp = np.random.multivariate_normal(mean, cov_x, num_particles ).T
y, yp = np.random.multivariate_normal(mean, cov_y, num_particles ).T

pm = np.column_stack((x, xp, y, yp))

distribution_plot(pm,[-40,40],[-6,6])

Notice that at the staring point, alpha_x,y is zero, 
 which means that the distribution is a typical ellipse.ex=sqrt(sx^2*sxp^2-sxxp^2)
sx0,sxp0,sy0,syp0 [4.812503660082882, 0.6527979769963828, 3.806466775496191, 0.8253303756159204]

# track the whole bunch for 100 turns and store their distribution
rm=[]
rm.append(pm) # store initial distribution
pnm=[]
ex0,ey0,sx0,sy0=distribution_parameters(pm)
pnm.append([0,ex0,ey0,sx0,sy0])
for i in range(1,1000):
    pm=track(TM,pm)
    rm.append(pm)
    ex,ey,sx,sy=distribution_parameters(pm)
    pnm.append([i+1,ex,ey,sx,sy])

# plot ex,ey,sx and sy in terms of i for pn

# Extract the data for plotting
turns = [entry[0] for entry in pnm]
parameters = {
    'Horizontal Emittance (ex)': [entry[1]*1000*1000 for entry in pnm],
    'Vertical Emittance (ey)': [entry[2]*1000*1000 for entry in pnm],
    'Horizontal Beam Size (sx)': [entry[3]*1000 for entry in pnm],
    'Vertical Beam Size (sy)': [entry[4]*1000 for entry in pnm]
}

# Create a figure and axis
fig, ax = plt.subplots(figsize=(12, 8))

# Plot each parameter
for param_name, param_values in parameters.items():
    linestyle = '--' if param_name in ['Horizontal Emittance (ex)', 'Vertical Emittance (ey)'] else '-'
    ax.plot(turns, param_values, label=param_name, linestyle=linestyle)

# Set labels and title
ax.set_xlabel('Turn Number (i)')
ax.set_ylabel('Parameter Value')
ax.set_title('Evolution of Emittance and Beam Size')

# Display legend and grid
ax.legend()
ax.grid(True)

# Show the plot
plt.show()

import matplotlib.pyplot as plt
import matplotlib.animation as animation

def distribution_animation(rn, file_output):
    fig, axs = plt.subplots(1, 3, figsize=(15, 5))

    n_frames = len(rn)
    xlim = [-30, 30]
    ylim = [-30, 30]

    def init():
        for ax in axs:
            ax.set_xlim(xlim)
            ax.set_ylim(ylim)

    def update(frame):
        for ax in axs:
            ax.clear()
            ax.set_xlim(xlim)
            ax.set_ylim(ylim)

        ex, ey, sx, sy = distribution_parameters(rn[frame])

        axs[0].scatter(rn[frame][:, 0]*1e3, rn[frame][:, 1]*1e3, s=5, color='b', alpha=0.5)
        axs[0].set_xlabel('x[mm]')
        axs[0].set_ylabel('xp[mrad]')
        axs[0].set_title('x-xp: {:.3f} pi.mm.mrad'.format(ex*1e6))

        axs[1].scatter(rn[frame][:, 2]*1e3, rn[frame][:, 3]*1e3, s=5, color='g', alpha=0.5)
        axs[1].set_xlabel('y[mm]')
        axs[1].set_ylabel('yp[mrad]')
        axs[1].set_title('y-yp: {:.3f} pi.mm.mrad'.format(ey*1e6))

        axs[2].scatter(rn[frame][:, 0]*1e3, rn[frame][:, 2]*1e3, s=5, color='r', alpha=0.5)
        axs[2].set_xlabel('x[mm]')
        axs[2].set_ylabel('y[mm]')
        axs[2].set_title('x - y')

        fig.suptitle('TURN {}'.format(frame), fontsize=16)

    ani = animation.FuncAnimation(fig, update, frames=n_frames, init_func=init, blit=False)
    ani.save(file_output, writer='pillow', fps=10)
    plt.show()

# Example usage
# Assuming you have defined the distribution_parameters function
# and have rn and file_output defined
# generate_particle_distribution_animation(rn, 'particle_distribution_n.gif')

distribution_animation(rm,"matched_run.gif")