Source code for ewoksid14.timepix4.t4_plot

import os
from typing import Optional
from typing import Union

import matplotlib.pyplot as plt
import numpy
import polars
from matplotlib.colors import LogNorm

from . import t4_camera
from . import t4_hist
from . import t4_plutils




def plot_hist2d(
    data: Union[polars.DataFrame, polars.LazyFrame, numpy.ndarray],
    x_pos: Optional[float] = None,
    y_pos: Optional[float] = None,
    title: str = "Spatial 2D histogram of events",
    filename: Optional[str] = None,
    show: bool = False,
):
    if not isinstance(data, numpy.ndarray):
        data = t4_hist.hist_2d(data)
    print(f"2D histogram min={data.min():,} max={data.max():,} sum={data.sum():,}")

    plt.figure()
    plt.imshow(data, origin="lower", norm=LogNorm(vmin=1), aspect="equal")

    if x_pos is not None:
        plt.axvline(x_pos, color="red", linewidth=1.5)
    if y_pos is not None:
        plt.axhline(y_pos, color="red", linewidth=1.5)

    plt.xlabel("X (colums)")
    plt.ylabel("Y (rows)")
    plt.tight_layout()
    plt.title(title)
    plt.colorbar(label="Counts (log scale)")
    _finalize_plot(filename=filename, show=show)





def plot_dxp_events(
    events: Union[polars.DataFrame, polars.LazyFrame],
    filename: Optional[str] = None,
    show: bool = False,
):
    """
    Plot 1D histogram of TOA for both DPX pixels (bot and top) on the same figure.
    """
    toa_bot = t4_plutils.collect(
        events.filter(t4_camera.is_dpx_bot_pixel).select(polars.col("toa48_ns"))
    )["toa48_ns"].to_numpy()
    toa_top = t4_plutils.collect(
        events.filter(t4_camera.is_dpx_top_pixel).select(polars.col("toa48_ns"))
    )["toa48_ns"].to_numpy()

    plt.figure()
    plt.scatter(toa_bot * 1e-9, [1] * toa_bot.size, s=1, color="blue", label="DPX Bot")
    plt.scatter(
        toa_top * 1e-9, [2] * toa_top.size, s=1, color="orange", label="DPX Top"
    )

    plt.xlabel("TOA (s)")
    plt.ylabel("Pixel")
    plt.yticks([1, 2], ["Bot", "Top"])
    plt.title("TOA scatter for DPX pixels")
    plt.legend()
    plt.tight_layout()

    print(f"Number of events: BOT={toa_bot.size:,}, TOP={toa_top.size:,}")

    _finalize_plot(filename=filename, show=show)



def _finalize_plot(filename: Optional[str] = None, show: bool = False):
    if show:
        plt.pause(0.1)
    if filename:
        dirname = os.path.dirname(filename)
        if dirname:
            os.makedirs(dirname, exist_ok=True)
        plt.savefig(filename, dpi=300, bbox_inches="tight")
        print(f"Saved {filename}")