plot_particle_tracking

Classes

ParticleTracking

ParticleTracking(
    run_number,
    z_pos="end",
    working_dir="",
    distribution_file="generator.in",
    run_file="photo_track.in",
    log_dir="log",
    log_file_name="plots_log.log",
    console_log=True,
    plots_dir="plots",
    n_R_bins=4,
    n_emittance_bins=4,
)

Bases: PlotPhaseSpace

Source code in plot_analysis/plot_particle_tracking.py

def __init__(self, run_number, z_pos="end", working_dir="", distribution_file="generator.in", run_file="photo_track.in", log_dir="log", log_file_name="plots_log.log", console_log=True, plots_dir="plots", n_R_bins=4, n_emittance_bins=4) -> None:
    super().__init__(run_number, z_pos, working_dir, distribution_file, run_file, log_dir, log_file_name, console_log, plots_dir)
    self.n_R_bins = n_R_bins
    self.n_emittance_bins = n_emittance_bins

    self.plots_dir = self.plots_dir.parents[0].joinpath(f"RUN_{run_number}_tracking")
    self.plots_dir.mkdir(exist_ok=True, parents=True)
    self.logger.info(f"Plots will be saved to: \n{self.plots_dir}")

    # create directories for position and emittance plots
    self.positions_dir = self.plots_dir.joinpath("position")
    self.positions_dir.mkdir(exist_ok=True, parents=True)
    self.emittance_dir = self.plots_dir.joinpath("emittance")
    self.emittance_dir.mkdir(exist_ok=True, parents=True)

    # get list of z points, where bunch properties are logged
    self.zpos_list = [int(f.name.split(".")[1]) for f in self.working_dir.iterdir() if f.is_file(
    ) and re.match(fr".*{self.run_file.stem}\.[0-9][0-9][0-9][0-9]\.{str(self.run_number).zfill(3)}", str(f))]

    # collect data at all z position
    self.data = self.collect_zpos_data()

Functions

collect_zpos_data

collect_zpos_data()

Collects the data for all zpos.

Returns:

Type	Description
DataFrame	pd.DataFrame: dataframe containing data from all the z positions.

Source code in plot_analysis/plot_particle_tracking.py

def collect_zpos_data(self) -> pd.DataFrame:
    """Collects the data for all zpos.

    Returns:
        pd.DataFrame: dataframe containing data from all the z positions.
    """
    data = pd.DataFrame([])
    particle_label = []
    emittance_label = []    

    for zpos in self.zpos_list:
        df = self.prepare_data(self.run_number, zpos)
        # make labels for particle R bin
        if data.shape[0] == 0:
            # for the first iteration, make the particle labels
            # calculate transversal distance from the center of the bunch
            R = np.sqrt(df["x"]**2 + df["y"]**2)
            # make the labelse
            particle_label = R // (R.max() / self.n_R_bins)
            # set R bin step for later use 
            self.R_bin_step = Decimal(str(self.get_step(R.max()/self.n_R_bins)))
        if len(emittance_label) == 0 and zpos > 0:
            # make the emittance labels and make sure the zpos == 0 is skipped (id does not make sense to look at the emittence there)
            # calculate the ,,emittance" of each particle
            x_px = np.sqrt(df["x"]**2 + df["px/pz"]**2)
            # calculate the emittance labels
            emittance_label = x_px // (x_px.max() / self.n_emittance_bins)
            # set emittance step for later use
            self.emittance_bin_step = Decimal(str(self.get_step(x_px.max()/self.n_emittance_bins)))

        df["R bin"] = particle_label
        df["emittance bin"] = emittance_label if len(emittance_label) != 0 else [np.NaN]*df.shape[0]
        data = pd.concat([data, df], ignore_index=True)
    return data

get_step

get_step(absolute_step)

Rounds the absolute_step so it is easily readable in the legend. If absolute_step > 1, simply rounds it to an integer. Otherwise looks for the first significant digit and rounds to that.

Parameters:

Name	Type	Description	Default
absolute_step	float	the absolute value of the step	required

Returns:

Name	Type	Description
float	float	rounded step

Source code in plot_analysis/plot_particle_tracking.py

def get_step(self, absolute_step:float) -> float:
    """Rounds the absolute_step so it is easily readable in the legend. If absolute_step > 1, simply rounds it to an integer. Otherwise looks for the first significant digit and rounds to that.

    Args:
        absolute_step (float): the absolute value of the step

    Returns:
        float: rounded step
    """
    # if larger than one, simply make in an int
    if absolute_step > 1:
        absolute_step = round(absolute_step, 0)
    # otherwise round it to the first significant digit
    else:
        # calculate the position of the first significat digit
        exponent = int(-math.floor(math.log10(abs(absolute_step))))
        # round the number to that digit
        absolute_step = round(absolute_step, exponent)
    return absolute_step

plot

plot()

Method to plot the positions and emittance. Main method to use.

Source code in plot_analysis/plot_particle_tracking.py

def plot(self):
    """Method to plot the positions and emittance. Main method to use.
    """
    self.logger.info("Plotting position development.")
    self.plot_positions()
    self.logger.info("Plotting emittance development.")
    self.plot_emittance()

plot_emittance

plot_emittance()

Runs the plot_emittance_at_zpos for each z position.

Source code in plot_analysis/plot_particle_tracking.py

def plot_emittance(self):
    """Runs the plot_emittance_at_zpos for each z position.
    """
    x_lim = self.data.loc[self.data["z_pos"] > 0]["x"].max()
    px_lim = self.data.loc[self.data["z_pos"] > 0]["px/pz"].max()
    for zpos in self.zpos_list:
        if zpos == 0:
            # there is no point in plotting at z_pos = 0, as the momentum is 0 there
            continue
        self.plot_emittance_scatter_at_zpos(zpos, x_lim=x_lim, px_lim=px_lim)

plot_emittance_scatter_at_zpos

plot_emittance_scatter_at_zpos(
    zpos, x_lim=None, px_lim=None, plot_autoscale=True
)

Plots scatter of the ,,emittance" of each particle at given zpos.

Parameters:

Name	Type	Description	Default
zpos	float	z position to be plotted	required
x_lim	float	Force sets the x_lim of the plot. If set to None, it scales automatically. Defaults to None.	None
px_lim	float	Force sets the y_lim of the plot. If set to None, it scales automatically.. Defaults to None.	None
plot_autoscale	bool	If True, makes sure that the autoscale plots are saved as well even if plot_lim is specified (in that case, two plots with different limits are saved). Defaults to True.	True

Source code in plot_analysis/plot_particle_tracking.py

def plot_emittance_scatter_at_zpos(self, zpos:float, x_lim=None, px_lim=None, plot_autoscale=True) -> None:
    """Plots scatter of the ,,emittance" of each particle at given zpos.

    Args:
        zpos (float): z position to be plotted
        x_lim (float, optional): Force sets the x_lim of the plot. If set to None, it scales automatically. Defaults to None.
        px_lim (float, optional): Force sets the y_lim of the plot. If set to None, it scales automatically.. Defaults to None.
        plot_autoscale (bool, optional): If True, makes sure that the autoscale plots are saved as well even if plot_lim is specified (in that case, two plots with different limits are saved). Defaults to True.
    """
    self.logger.info(f"Plotting emittance scatter at z={zpos/100} m.")

    df = self.data.loc[self.data["z_pos"] == zpos]

    fig = plt.figure(figsize=(6, 6))
    # create the grid
    gs = fig.add_gridspec(2, 2,  width_ratios=(3, 1), height_ratios=(1, 3),
                  left=0.1, right=0.9, bottom=0.1, top=0.9,
                  wspace=0.05, hspace=0.05)

    # add scatter to the grid
    ax = fig.add_subplot(gs[1, 0])
    for bin_n in range(self.n_emittance_bins):
        dff = df.loc[df["emittance bin"] == bin_n]
        ax.scatter(dff["x"], dff["px/pz"], s=1)

    # add histograms to the grid
    ax_histx = fig.add_subplot(gs[0, 0], sharex=ax)
    ax_histy = fig.add_subplot(gs[1, 1], sharey=ax)
    ax_histx.tick_params(axis="x", labelbottom=False)
    ax_histy.tick_params(axis="y", labelleft=False)

    v_counts = np.vstack(
        [np.histogram(df.loc[df["emittance bin"] == b]["px/pz"], bins=50)[0] for b in range(self.n_R_bins)]
    ).transpose()
    v_bins = np.vstack(
        [np.histogram(df.loc[df["emittance bin"] == b]["px/pz"], bins=50)[1][:-1] for b in range(self.n_R_bins)]
    ).transpose()
    ax_histy.hist(v_bins, bins=v_bins.shape[0], weights = v_counts, orientation='horizontal', stacked=True)

    h_counts = np.vstack(
        [np.histogram(df.loc[df["emittance bin"] == b]["x"], bins=50)[0] for b in range(self.n_R_bins)]
    ).transpose()
    h_bins = np.vstack(
        [np.histogram(df.loc[df["emittance bin"] == b]["x"], bins=50)[1][:-1] for b in range(self.n_R_bins)]
    ).transpose()
    ax_histx.hist(h_bins, bins=h_bins.shape[0], weights = h_counts, stacked=True)

    # make inner legend with classic hist markers instead of the scatter ones
    ax.plot([], [], " ", label=f"zpos = {zpos/100} m")
    for i in range(self.n_emittance_bins):
        label = f"${i*self.emittance_bin_step} \leq x\cdot p_x $" + (f"$< {(i+1)*self.emittance_bin_step}$" if i != self.n_R_bins-1 else "")
        ax.scatter([], [], c=self.default_colors[i], label= label, marker="s")
    ax.legend()

    ax.set_xlabel("$x$ [mm]")
    ax.set_ylabel("$p_x/p_z$ [mrad]")

    ax_histx.set_ylabel("N")
    ax_histy.set_xlabel("N")

    # save the autoscaled plot  
    if plot_autoscale or (x_lim is None and px_lim is None):
        # if plot_lim is specified, save the autoscale plots to a separate dir
        autoscale_dir = self.emittance_dir.joinpath("autoscale") if not (x_lim is None and px_lim is None) else self.emittance_dir
        autoscale_dir.mkdir(exist_ok=True, parents=True)
        fig_name = autoscale_dir.joinpath(f"T_view_{zpos}.png")
        fig.savefig(fig_name, dpi=300, bbox_inches="tight")

    # force set the limits if specified
    if px_lim is not None:
        ax.set_ylim(-abs(px_lim), abs(px_lim))
    if x_lim is not None:
        ax.set_xlim(-abs(x_lim), abs(x_lim))
    if x_lim is None and px_lim is None:
        plt.close()
        return

    fig_name = self.emittance_dir.joinpath(f"emittance_zpos_{zpos}.png")
    fig.savefig(fig_name, dpi=300, bbox_inches="tight")
    plt.close()

plot_position_at_zpos

plot_position_at_zpos(
    zpos, plot_lim=None, plot_autoscale=True
)

Plots scatter of transversal particle position at givin zpos.

Parameters:

Name	Type	Description	Default
zpos	float	z position, to be plotted	required
plot_lim	float	Force set the limits of the plot, if set to None, it scales automatically. Defaults to None.	None
plot_autoscale	bool	If True, makes sure that the autoscale plots are saved as well even if plot_lim is specified (in that case, two plots with different limits are saved). Defaults to True.	True

Source code in plot_analysis/plot_particle_tracking.py

def plot_position_at_zpos(self, zpos:float, plot_lim=None, plot_autoscale=True) -> None:
    """Plots scatter of transversal particle position at givin zpos.

    Args:
        zpos (float): z position, to be plotted
        plot_lim (float, optional): Force set the limits of the plot, if set to None, it scales automatically. Defaults to None.
        plot_autoscale (bool, optional): If True, makes sure that the autoscale plots are saved as well even if plot_lim is specified (in that case, two plots with different limits are saved). Defaults to True.
    """
    self.logger.info(f"Plotting particle position at z={zpos/100} m.")

    df = self.data.loc[self.data["z_pos"] == zpos]

    fig = plt.figure(figsize=(6, 6))
    # create the grid
    gs = fig.add_gridspec(2, 2,  width_ratios=(3, 1), height_ratios=(1, 3),
                  left=0.1, right=0.9, bottom=0.1, top=0.9,
                  wspace=0.05, hspace=0.05)

    # add scatter to the grid
    ax = fig.add_subplot(gs[1, 0])
    for bin_n in range(self.n_R_bins):
        dff = df.loc[df["R bin"] == bin_n]
        ax.scatter(dff["x"], dff["y"], s=1)

    # add histograms to the grid
    ax_histx = fig.add_subplot(gs[0, 0], sharex=ax)
    ax_histy = fig.add_subplot(gs[1, 1], sharey=ax)
    ax_histx.tick_params(axis="x", labelbottom=False)
    ax_histy.tick_params(axis="y", labelleft=False)

    # calculate the content of the bins based on the R bin of the particle
    v_counts = np.vstack(
        [np.histogram(df.loc[df["R bin"] == b]["y"], bins=50)[0] for b in range(self.n_R_bins)]
    ).transpose()
    v_bins = np.vstack(
        [np.histogram(df.loc[df["R bin"] == b]["y"], bins=50)[1][:-1] for b in range(self.n_R_bins)]
    ).transpose()
    ax_histy.hist(v_bins, bins=v_bins.shape[0], weights = v_counts, orientation='horizontal', stacked=True)

    h_counts = np.vstack(
        [np.histogram(df.loc[df["R bin"] == b]["x"], bins=50)[0] for b in range(self.n_R_bins)]
    ).transpose()
    h_bins = np.vstack(
        [np.histogram(df.loc[df["R bin"] == b]["x"], bins=50)[1][:-1] for b in range(self.n_R_bins)]
    ).transpose()
    ax_histx.hist(h_bins, bins=h_bins.shape[0], weights = h_counts, stacked=True)

    # make inner legend with classic hist markers instead of the scatter ones
    ax.plot([], [], " ", label=f"zpos = {zpos/100} m")
    for i in range(self.n_R_bins):
        label = f"${i*self.R_bin_step} \leq R$" + (f"$< {(i+1)*self.R_bin_step}$" if i != self.n_R_bins-1 else "")
        ax.scatter([], [], c=self.default_colors[i], label= label, marker="s")
    ax.legend(loc="upper right")

    ax.set_xlabel("$\Delta x$ [mm]")
    ax.set_ylabel("$\Delta y$ [mm]")

    ax_histx.set_ylabel("N")
    ax_histy.set_xlabel("N")

    # save the autoscaled plot
    if plot_autoscale or plot_lim is None:
        # if plot_lim is specified, save the autoscale plots to a separate dir
        autoscale_dir = self.positions_dir.joinpath("autoscale") if plot_lim is not None else self.positions_dir
        autoscale_dir.mkdir(exist_ok=True, parents=True)
        fig_name = autoscale_dir.joinpath(f"T_view_{zpos}.png")
        fig.savefig(fig_name, dpi=300, bbox_inches="tight")

    # force set the plot limits if specified
    if plot_lim is not None:
        ax.set_ylim(-abs(plot_lim), abs(plot_lim))
        ax.set_xlim(-abs(plot_lim), abs(plot_lim))

        fig_name = self.positions_dir.joinpath(f"T_view_{zpos}.png")
        fig.savefig(fig_name, dpi=300, bbox_inches="tight")

    plt.close()

plot_positions

plot_positions()

Runs the plot_position_at_zpos for each z position.

Source code in plot_analysis/plot_particle_tracking.py

def plot_positions(self) -> None:
    """Runs the plot_position_at_zpos for each z position. 
    """
    plot_lim = max(self.data["x"].max(), self.data["y"].max())
    for zpos in self.zpos_list:
        self.plot_position_at_zpos(zpos, plot_lim=plot_lim)

plot_transverse_particle_density

plot_transverse_particle_density()

Runs the plot_transvese_particle_density_at_zpos for each z position.

Source code in plot_analysis/plot_particle_tracking.py

def plot_transverse_particle_density(self) -> None:
    """Runs the plot_transvese_particle_density_at_zpos for each z position.
    """
    x_lim = self.data["x"].max()
    y_lim = self.data["y"].max()
    for zpos in self.zpos_list:
        self.plot_transverse_particle_density_at_zpos(zpos, x_lim=x_lim, y_lim=y_lim)

plot_transverse_particle_density_at_zpos

plot_transverse_particle_density_at_zpos(
    zpos, x_lim=None, y_lim=None, plot_autoscale=True
)

Plots transverse particle density (heatmap) at given zpos.

Parameters:

Name	Type	Description	Default
zpos	int	z position, at which to plot	required
x_lim	float	Limit the x axis of the heatmap. If not specified (set to None) the x axis is autoscaled. Defaults to None.	None
y_lim	_type_	Limit the y axis of the heatmap. If not specified (set to None) the y axis is autoscaled.. Defaults to None.	None
plot_autoscale	bool	If True, makes sure that the autoscale plots are saved as well even if plot_lim is specified (in that case, two plots with different limits are saved). Defaults to True.	True

Source code in plot_analysis/plot_particle_tracking.py

def plot_transverse_particle_density_at_zpos(self, zpos:int, x_lim=None, y_lim=None, plot_autoscale=True) -> None:
    """Plots transverse particle density (heatmap) at given zpos.

    Args:
        zpos (int): z position, at which to plot
        x_lim (float, optional): Limit the x axis of the heatmap. If not specified (set to None) the x axis is autoscaled. Defaults to None.
        y_lim (_type_, optional): Limit the y axis of the heatmap. If not specified (set to None) the y axis is autoscaled.. Defaults to None.
        plot_autoscale (bool, optional): If True, makes sure that the autoscale plots are saved as well even if plot_lim is specified (in that case, two plots with different limits are saved). Defaults to True.
    """
    # get the plot from its build function
    fig, gs, ax, ax_histy, ax_histx = self._transverse_particle_density_plot(self.z_pos)

    # set the scaling
    if plot_autoscale or (x_lim is None and y_lim is None):
        autoscale_dir = self.emittance_dir.joinpath("autoscale") if plot_autoscale else self.emittance_dir
        autoscale_dir.mkdir(exist_ok=True, parents=True)
        fig_name = autoscale_dir.joinpath(f"T_view_{zpos}.png")
        fig.savefig(fig_name, dpi=300, bbox_inches="tight")

    if y_lim is not None:
        ax.set_ylim(-abs(y_lim), abs(y_lim))
    if x_lim is not None:
        ax.set_xlim(-abs(x_lim), abs(x_lim))
    if x_lim is None and y_lim is None:
        plt.close()
        return

    fig_name = self.emittance_dir.joinpath(f"emittance_zpos_{zpos}.png")
    fig.savefig(fig_name, dpi=300, bbox_inches="tight")
    plt.close()

save_data

save_data()

Function to save the data to an excel file.

Source code in plot_analysis/plot_particle_tracking.py

def save_data(self):
    """Function to save the data to an excel file.
    """
    # is it neccessary?
    save_file = self.plots_dir.joinpath("data.xlsx")
    self.logger.info(f"Saving data to an excel file:\n{save_file}")
    self.data.to_excel(save_file)

---

Last update: October 31, 2023
Created: October 31, 2023