draw_tracks.py

Go to the documentation of this file.

#/// script
#dependencies = [
#"numpy",
#"matplotlib"
#]
#///
 
""
    "
    Draw a random subset of 3D straight
    - line tracks from a detector CSV file.
 
      Each row encodes a track as x = xA + xB * z,
                                  y = yA
                                      + yB
                                            * z with columns : 0
    : x 1
    : y 2
    : z 3
    : xA 4
    : xB 5
    : yA 6
    : yB
 
          Lines are drawn over the full z range present in the file.The measured hit
      point(x, y, z) is also marked.
"""
 
import argparse
import sys
 
import matplotlib.pyplot as plt
import numpy as np
 
# ── column indices(same as train.py) ─────────────────────────────────────────
iX,  iY,  iZ       = 0, 1, 2
iXA, iXB, iYA, iYB = 3, 4, 5, 6
 
DEFAULT_N = 50
 
 
def main():
    parser = argparse.ArgumentParser(
        description="Draw random 3D tracks from a detector CSV file."
    )
    parser.add_argument("csv",           help="path to detector CSV file")
    parser.add_argument("-n", "--count", default=DEFAULT_N, type=int,
                        help=f"number of random tracks to draw (default {DEFAULT_N})")
    parser.add_argument("--seed",        default=None, type=int,
                        help="random seed for reproducibility")
    args = parser.parse_args()
 
    data = np.loadtxt(args.csv, delimiter=",", skiprows=1, dtype=np.float64)
    if data.ndim == 1:
        data = data[np.newaxis, :]
 
    rng = np.random.default_rng(args.seed)
    n   = min(args.count, len(data))
    idx = rng.choice(len(data), size=n, replace=False)
    rows = data[idx]
 
    z_max = data[:, iZ].max()
    z_lo  = 0.0
    z_hi  = z_max + 0.1 * max(z_max, 1.0)
    zz   = np.linspace(z_lo, z_hi, 50)
 
    fig = plt.figure(figsize=(10, 7))
    ax  = fig.add_subplot(111, projection="3d")
 
    for row in rows:
        xA, xB = row[iXA], row[iXB]
        yA, yB = row[iYA], row[iYB]
        xs = xA + xB * zz
        ys = yA + yB * zz
        ax.plot(zz, xs, ys, color="steelblue", alpha=0.4, linewidth=0.8)
 
#measured hit points
    ax.scatter(rows[:, iZ], rows[:, iX], rows[:, iY],
               color="tomato", s=12, zorder=5, label="measured hits")
 
    ax.set_xlabel("Z")
    ax.set_ylabel("X")
    ax.set_zlabel("Y")
    ax.set_title(f"{n} random tracks  —  {args.csv}")
    ax.legend()
    plt.tight_layout()
    plt.show()
 
 
if __name__ == "__main__":
    main()