train.py

Go to the documentation of this file.

#/// script
#dependencies = [
#"torch==2.6.0",
#"numpy"
#]
#
#[tool.uv.sources]
#torch = {index = "pytorch" }
#
#[[tool.uv.index]]
#name = "pytorch"
#url = "https://download.pytorch.org/whl/cu124"
#///
 
"""
    BM@N alignment routine
    BM@N experiment at NICA complex, JINR, 2025
 
    Department: Math & Soft Group of HEP lab
    Author:     Igor Polev, polev@jinr.ru
 
    MSE minimization of alignment problem using NN training techniques.
    For each alignable detector element we prepare a custom single-layer NN model
    and train it to get trained NN weights as a solution for alignment parameters.
    Custom layer consists of measuremets model for defined experiment setup.
    Datasets are prepared by expdataNN.C from original hits and tracks
    alignment data in form of separate CSV files for each detector element:
 
    inputs  (5): z, xA, xB, yA, yB   — columns [2..6]
    outputs (2): x, y                — columns [0..1]
 
    where xA, xB, yA, yB are coefficients of straight-line fit:
    
    x = xA + xB*z, y = yA + yB*z
 
    This fit is computed from non-aligned data using ordinary linear regression
    during data export procedure (see expdataNN.C).
"""
 
import argparse
import glob
import json
import os
import re
import sys
 
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
 
# ── defaults ──────────────────────────────────────────────────────────────────
DEFAULT_GLOB    = "/home/igor/DATA/bmn/nn/data_det*.csv"
DEFAULT_EPOCHS  = 100
DEFAULT_LR      = 1e-3
DEFAULT_BATCH   = 256
DEFAULT_OUT_DIR = "models"
 
# ── dataset layout ────────────────────────────────────────────────────────────
iX,  iY,  iZ       = 0, 1, 2
iXA, iXB, iYA, iYB = 3, 4, 5, 6
 
INPUT_COLS  = [iZ, iXA, iYA, iXB, iYB]
OUTPUT_COLS = [iX, iY]
 
# ── NN definition ─────────────────────────────────────────────────────────────
#
#Custom layer to mimic alignment task as a neural network
#
#Input : z, xA, xB, yA, yB(shape : [ batch, 5 ])
#Output : x, y(shape : [ batch, 2 ])
#
#NN parameters = alignment parameners:
#shift = [dx, dy, dz]
#
#layers math is measurements model accounting for (mis) alignment params
 
class AlignLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.shiftV = nn.Parameter(torch.zeros(2))
        self.shiftZ = nn.Parameter(torch.zeros(1))
 
    def forward(self, batch: torch.Tensor) -> torch.Tensor:
        A = batch[:, [iXA - 2, iYA - 2]]
        B = batch[:, [iXB - 2, iYB - 2]]
        z = batch[:, [iZ - 2]] - self.shiftZ
 
        return A + B * z + self.shiftV
 
def load_dataset(path: str) -> TensorDataset:
    raw = np.loadtxt(path, delimiter=",", skiprows=1, dtype=np.float32)
    return TensorDataset(
        torch.from_numpy(raw[:, INPUT_COLS]),
        torch.from_numpy(raw[:, OUTPUT_COLS])
    )
 
def train_one(
    dataset:   TensorDataset,
    det_id:    str,
    epochs:    int,
    lr:        float,
    batch:     int,
    out_dir:   str,
    device:    torch.device,
) -> tuple[float, float, float]:
    """Train and return (dx, dy, dz) weights."""
    loader  = DataLoader(dataset, batch_size=batch, shuffle=True)
    model   = AlignLayer().to(device)
    opt     = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
 
    for epoch in range(1, epochs + 1):
        model.train()
        total_loss = 0.0
        for X_batch, Y_batch in loader:
            X_batch = X_batch.to(device)
            Y_batch = Y_batch.to(device)
            opt.zero_grad()
            loss = loss_fn(model(X_batch), Y_batch)
            loss.backward()
            opt.step()
            total_loss += loss.item() * len(X_batch)
 
        if epoch % max(1, epochs // 10) == 0 or epoch == epochs:
            mse = total_loss / len(dataset)
            with torch.no_grad():
                dx, dy = model.shiftV.cpu().tolist()
                dz     = model.shiftZ.cpu().item()
            print(f"  [{det_id}] epoch {epoch:>{len(str(epochs))}}/{epochs}  MSE={mse:.6e}  dx={dx:+.6f}  dy={dy:+.6f}  dz={dz:+.6f}")
 
    os.makedirs(out_dir, exist_ok=True)
    save_path = os.path.join(out_dir, f"net_{det_id}.pt")
    torch.save(model.state_dict(), save_path)
    print(f"  [{det_id}] saved -> {save_path}")
 
    with torch.no_grad():
        dx, dy = model.shiftV.cpu().tolist()
        dz     = model.shiftZ.cpu().item()
    return dx, dy, dz
 
def main():
    parser = argparse.ArgumentParser(
        description="Train one linear NN per detector CSV file."
    )
    parser.add_argument("--data",     default=DEFAULT_GLOB,    help="glob pattern for CSV files")
    parser.add_argument("--epochs",   default=DEFAULT_EPOCHS,  type=int)
    parser.add_argument("--lr",       default=DEFAULT_LR,      type=float)
    parser.add_argument("--batch",    default=DEFAULT_BATCH,   type=int)
    parser.add_argument("--outdir",   default=DEFAULT_OUT_DIR, help="directory to save trained models")
    parser.add_argument("--solution", default=None,            help="path to save solution JSON (default: <outdir>/solution.json)")
    args = parser.parse_args()
 
    csv_files = sorted(glob.glob(args.data))
    if not csv_files:
        print(f"No CSV files matched: {args.data}", file=sys.stderr)
        sys.exit(1)
 
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    print(f"Found {len(csv_files)} detector file(s)\n")
 
    corrections = []
    for path in csv_files:
        stem   = os.path.splitext(os.path.basename(path))[0]  # e.g. "data_det12345"
        m      = re.search(r"(\d+)$", stem)
        det_id = int(m.group(1)) if m else stem
        print(f"==> {path}  (detector {det_id})")
        dataset = load_dataset(path)
        print(f"    samples: {len(dataset)}")
        dx, dy, dz = train_one(dataset, stem, args.epochs, args.lr, args.batch, args.outdir, device)
        corrections.append({"detector_id": det_id, "values": [dx, dy, dz]})
        print()
 
    solution = {
        "CorrectionValues": ["dX", "dY", "dZ"],
        "CorrectionsPerDetector": corrections,
        "DetectorElements": len(corrections),
    }
    solution_path = args.solution or os.path.join(args.outdir, "solution.json")
    os.makedirs(os.path.dirname(os.path.abspath(solution_path)), exist_ok=True)
    with open(solution_path, "w") as f:
        json.dump(solution, f, indent=4)
    print(f"Solution saved -> {solution_path}")
    print("All done.")
 
if __name__ == "__main__":
    main()