Modules

load_geometry(file_path, tree_name)

Load geometry from a ROOT file into a pandas DataFrame.

Parameters:

Name	Type	Description	Default
file_path	str	The path to the ROOT file.	required
tree_name	str	The name of the tree to load.	required

Returns:

Type	Description
DataFrame	pd.DataFrame: A pandas DataFrame containing the loaded geometry.

Raises:

Type	Description
Exception	If coordinate branches have not been set before loading geometry.

Source code in pygeosimplify/io/geo_handler.py

def load_geometry(file_path: str, tree_name: str) -> pd.DataFrame:
    """
    Load geometry from a ROOT file into a pandas DataFrame.

    Args:
        file_path (str): The path to the ROOT file.
        tree_name (str): The name of the tree to load.

    Returns:
        pd.DataFrame: A pandas DataFrame containing the loaded geometry.

    Raises:
        Exception: If coordinate branches have not been set before loading geometry.
    """
    if config.coordinate_branch_names == {}:
        raise Exception(
            "Coordinate branches have not been set. Please set coordinate branches before loading geometry."
        )
    # Open root tree with uprot
    tree = uproot.open(f"{file_path}:{tree_name}")
    # Convert the tree to a pandas dataframe
    df = tree_to_df(tree)
    # Check whether the tree contains all required branches
    check_geo_consistency(df)

    return df

plot_geometry(df, ax=None, layer_list=None, eta_range=None, phi_range=None, axis_labels=None, color=None, unit_scale=1, cell_energy_col=None, unit_scale_energy=1, energy_label='Cell Energy', color_map='gist_heat_r')

Plot the geometry based on the provided DataFrame.

Parameters:

Name	Type	Description	Default
df	DataFrame	The DataFrame containing the geometry data.	required
ax	Axes3D	The 3D axes to plot on. If not provided, a new figure and axes will be created.	None
layer_list	list[int]	The list of layers to consider. If not provided, all layers will be considered.	None
eta_range	list	The range of eta values to filter the data. If not provided, the default range is [-5, 5].	None
phi_range	list	The range of phi values to filter the data. If not provided, the default range is [0, np.pi].	None
axis_labels	list	The labels for the x, y, and z axes. If not provided, the default labels are ["x", "y", "z"].	None
color	str	The color to use for the cells. If not provided, colors will be automatically assigned based on the layers.	None
unit_scale	float	The scale factor for the unit of measurement. Default is 1.	1
cell_energy_col	str	The name of the column containing the cell energy values. If provided, the cells will be colored based on the energy values.	None
unit_scale_energy	float	The scale factor for the unit of measurement of the cell energy. Default is 1.	1
energy_label	str	The label for the colorbar when cell energy is used. Default is "Cell Energy".	'Cell Energy'
color_map	str	The colormap to use when coloring the cells based on energy values. Default is "gist_heat_r".	'gist_heat_r'

Returns:

Name	Type	Description
Axes3D	Axes3D	The 3D axes object containing the plot.

Source code in pygeosimplify/vis/geo.py

def plot_geometry(  # noqa: C901
    df: pd.DataFrame,
    ax: Union[Axes3D, None] = None,
    layer_list: list[int] | None = None,
    eta_range: list | None = None,
    phi_range: list | None = None,
    axis_labels: list | None = None,
    color: str | None = None,
    unit_scale: float = 1,
    cell_energy_col: str | None = None,
    unit_scale_energy: float = 1,
    energy_label: str = "Cell Energy",
    color_map: str = "gist_heat_r",
) -> Axes3D:
    """
    Plot the geometry based on the provided DataFrame.

    Parameters:
        df (pd.DataFrame): The DataFrame containing the geometry data.
        ax (Axes3D, optional): The 3D axes to plot on. If not provided, a new figure and axes will be created.
        layer_list (list[int], optional): The list of layers to consider. If not provided, all layers will be considered.
        eta_range (list, optional): The range of eta values to filter the data. If not provided, the default range is [-5, 5].
        phi_range (list, optional): The range of phi values to filter the data. If not provided, the default range is [0, np.pi].
        axis_labels (list, optional): The labels for the x, y, and z axes. If not provided, the default labels are ["x", "y", "z"].
        color (str, optional): The color to use for the cells. If not provided, colors will be automatically assigned based on the layers.
        unit_scale (float, optional): The scale factor for the unit of measurement. Default is 1.
        cell_energy_col (str, optional): The name of the column containing the cell energy values. If provided, the cells will be colored based on the energy values.
        unit_scale_energy (float, optional): The scale factor for the unit of measurement of the cell energy. Default is 1.
        energy_label (str, optional): The label for the colorbar when cell energy is used. Default is "Cell Energy".
        color_map (str, optional): The colormap to use when coloring the cells based on energy values. Default is "gist_heat_r".

    Returns:
        Axes3D: The 3D axes object containing the plot.
    """

    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")

    if eta_range is None:
        eta_range = [-5, 5]
    if phi_range is None:
        phi_range = [0, np.pi]
    if axis_labels is None:
        axis_labels = ["x", "y", "z"]

    # If no layer list is provided consider all all layers
    if layer_list is None:
        layer_list = list(df["layer"].unique())

    # Filter for layer list
    df = df[df["layer"].isin(layer_list)]

    # Filter for eta and phi range
    df = filter_df_eta_phi(df, eta_range, phi_range)

    # Create a visual cell scene
    vis = CellScene()

    if not cell_energy_col:
        # Create a color dict mapping a layer to a color
        layer_color_dict = dict(zip(layer_list, get_colors(len(layer_list), rng=0), strict=False))
        # If color is specifically provided, override the color dict
        if color:
            layer_color_dict = dict.fromkeys(layer_list, color)
        add_cells_to_scene(
            df=df,
            scene=vis,
            unit_scale=unit_scale,
            layer_color_dict=layer_color_dict,
        )
    else:
        # Make sure the energy column exists
        if cell_energy_col not in df.columns:
            raise ValueError(f"Column {cell_energy_col} not found in DataFrame")
        # Make sure the energy column is not empty and not always 0
        if df[cell_energy_col].empty or df[cell_energy_col].eq(0).all():
            raise ValueError(f"Column {cell_energy_col} is empty or always 0")

        # Create a color map mapping cell energy to a color
        vmin = df[cell_energy_col].min() * unit_scale_energy
        vmax = df[cell_energy_col].max() * unit_scale_energy
        norm = mcolors.LogNorm(vmin * 0.1, vmax)

        add_cells_to_scene(
            df=df,
            scene=vis,
            unit_scale=unit_scale,
            unit_scale_energy=unit_scale_energy,
            colormap=plt.get_cmap(color_map),
            norm=norm,
        )

    vis.plot(ax=ax, axis_labels=axis_labels)

    if cell_energy_col:
        mappable = plt.cm.ScalarMappable(norm=norm, cmap=plt.get_cmap(color_map))
        cbar = plt.colorbar(mappable, ax=ax, fraction=0.035, pad=0.15)
        cbar.set_label(energy_label)

    # Regularize x,y limits so that limits are identical for x and y (to avoid distortions)
    minMaxX = vis.min_max_cell_list_extent(0)
    minMaxY = vis.min_max_cell_list_extent(1)
    minMax: tuple[float, float] = (min(minMaxX[0], minMaxY[0]), max(minMaxX[1], minMaxY[1]))

    ax.set_xlim(minMax)
    ax.set_ylim(minMax)

    return ax