Tensor Network for Anomaly Detection in Latent Space of Proton-Proton Collision Events at the LHC#

Implementation of the anomaly detection pipeline from the paper: 10.1088/2632-2153/ae0243.

Install tn4ml directly from GitHub:

git clone https://github.com/bsc-quantic/tn4ml.git

Install an additional package for data handling:

h5py

(optional) Download dataset#

The dimensionality of the dataset is reduced by passing it through autoencoder. If you are interested more in the autoencoder’s architecture, please refer to [*].

Reduced dataset can be downloaded from Zenodo:

record/7673769

Description of filenames:

latentrep_QCD_sig.h5: train dataset (QCD - background)
latentrep_QCD_sig_testclustering.h5: test dataset (QCD - background)
latentrep_RSGraviton_WW_NA_35.h5: test dataset (Signal $mathrm{NA G rightarrow WW}$)

Or it can be downloaded directly in the pipeline.py.

Run training and evaluation pipeline#

Data Parameters

save_dir (str): Path to directory for saving results (default = "results/")
load_dir (str): Path to directory for loading the data
feature_range (list): Feature range for scaling (default = [0, 1])
seed (int): Seed for random number generator
standardization (str): Standardization of data ("yes" or "no", default = "yes")
minmax (str): Minmax scaling of data ("yes" or "no", default = "yes")
embedding (str): Embedding type for input data (e.g., "legendre_4", "fourier_2", "hermite_3")
test_size (int): Number of samples for testing
train_size (int): Number of samples for training
signal_name (str): Name of signal dataset (default = "RSGraviton_WW_NA_35")
latent (int): Latent space dimension

MPS Parameters

bond_dim (int): Bond dimension of MPS (default = 5)
initializer (str): Type of MPS initialization
shape_method (str): Method for distributing bond dimensions (default = "even")

Training Parameters

lr (float): Learning rate (default = 1e-3)
min_delta (float): Minimum improvement required for early stopping (default = 0)
patience (int): Number of epochs with no improvement before early stopping (default = 20)
epochs (int): Maximum number of training epochs (default = 100)
batch_size (int): Number of samples per training batch (default = 32)
run (int): Number of training repetitions with different seeds

python pipeline.py -save_dir results/ \
                   -load_dir QML_paper_data \
                   -feature_range 0 1 \
                   -minmax yes \
                   -embedding laguerre_2 \
                   -test_size 5000 \
                   -train_size 10000 \
                   -bond_dim 8 \
                   -initializer unitary \
                   -lr 0.001 \
                   -patience 25 \
                   -epochs 100 \
                   -batch_size 128 \
                   -run 1 \
                   -latent 4

Run evaluation only#

Data Parameters

save_dir (str): Path to directory for saving results (default = "results/")
load_dir (str): Path to directory for loading the data
feature_range (list): Feature range for scaling (default = [0, 1])
seed (int): Seed for random number generator
standardization (str): Standardization of data ("yes" or "no", default = "yes")
minmax (str): Minmax scaling of data ("yes" or "no", default = "yes")
embedding (str): Embedding type for input data (e.g., "legendre_4", "fourier_2", "hermite_3")
test_size (int): Number of samples for testing
signal_name (str): Name of signal dataset (default = "RSGraviton_WW_NA_35", options: "RSGraviton_WW_BR_15", "AtoHZ_to_ZZZ_35")
latent (int): Latent space dimension

MPS Parameters

bond_dim (int): Bond dimension of MPS (default = 5)
initializer (str): Type of MPS initialization

Training Parameters

batch_size (int): Number of samples per training batch (default = 32)
run (int): Number of training repetitions with different seeds

python evaluation.py -save_dir results/ \
                   -load_dir QML_paper_data \
                   -feature_range 0 1 \
                   -minmax yes \
                   -embedding laguerre_2 \
                   -test_size 5000 \
                   -bond_dim 8 \
                   -initializer unitary \
                   -batch_size 128 \
                   -run 1 \
                   -latent 4

Additional Example Scripts#

Below are the Python scripts used in this example. You can view their contents directly on this page, or access them on GitHub.

evaluation.py: evaluation.py View on GitHub

evaluation.py#

import argparse
import os

import h5py
import jax
import jax.numpy as jnp
import joblib
import numpy as np
from utils import (
    Colors,
    _ensure_data_exists,
    calc_fidelity_batch,
    load_test_data,
)

import tn4ml
from tn4ml.embeddings import (
    Embedding,
    FourierEmbedding,
    HermiteEmbedding,
    LaguerreEmbedding,
    LegendreEmbedding,
)
from tn4ml.initializers import gramschmidt, rand_unitary
from tn4ml.models.model import load_model

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="read arguments for training of TN model"
    )
    parser.add_argument(
        "-save_dir",
        dest="save_dir",
        type=str,
        help="path to directory for saving results",
        default="results/",
    )
    parser.add_argument(
        "-load_dir",
        dest="load_dir",
        type=str,
        help="path to directory for loading the data",
    )

    # data params
    parser.add_argument(
        "-feature_range",
        dest="feature_range",
        type=float,
        nargs=2,
        default=[0, 1],
        help="Feature range for scaling",
    )
    parser.add_argument(
        "-seed", dest="seed", type=int, help="Seed for random number generator"
    )
    parser.add_argument(
        "-standardization",
        dest="standardization",
        type=str,
        default="yes",
        choices=["yes", "no"],
        help="Standardization of data",
    )
    parser.add_argument(
        "-minmax",
        dest="minmax",
        type=str,
        default="yes",
        choices=["yes", "no"],
        help="Minmax scaling of data",
    )
    parser.add_argument(
        "-embedding", dest="embedding", type=str, help="Embedding type for input data"
    )
    parser.add_argument("-test_size", dest="test_size", type=int, help="Test size")
    parser.add_argument(
        "-signal_name",
        dest="signal_name",
        type=str,
        default="RSGraviton_WW_NA_35",
        help="Name of signal",
    )

    # MPS params
    parser.add_argument(
        "-bond_dim", dest="bond_dim", type=int, default=5, help="Bond dimension"
    )
    parser.add_argument(
        "-initializer", dest="initializer", type=str, help="Type of MPS initialization"
    )

    # testing params
    parser.add_argument("-batch_size", dest="batch_size", type=int, default=32)
    parser.add_argument(
        "-run", dest="run", type=int, help="Number of training repetitions"
    )

    parser.add_argument(
        "-latent", dest="latent", type=int, help="Latent space dimension"
    )

    args = parser.parse_args()
    params = vars(args)

    # Get paths to all data files, downloading them if necessary
    print(
        Colors.YELLOW.value + "Checking data folder..." + Colors.RESET.value + "\n",
        end="",
    )

    _ensure_data_exists(args.load_dir, args.latent)

    print(Colors.BLUE.value + "Importing data... " + Colors.RESET.value + "\n", end="")

    if args.standardization == "yes":
        save_dir = (
            args.save_dir
            + "/"
            + args.initializer
            + "/10k_standard/"
            + str(args.embedding)
            + "/lat"
            + str(args.latent)
            + "/bond"
            + str(args.bond_dim)
            + "/run_"
            + str(args.run)
        )
    elif args.minmax == "yes":
        if tuple(args.feature_range) == (-1, 1):
            save_dir = (
                args.save_dir
                + "/"
                + args.initializer
                + "/10k_minmax-11/"
                + str(args.embedding)
                + "/lat"
                + str(args.latent)
                + "/bond"
                + str(args.bond_dim)
                + "/run_"
                + str(args.run)
            )
        else:
            save_dir = (
                args.save_dir
                + "/"
                + args.initializer
                + "/10k_minmax01/"
                + str(args.embedding)
                + "/lat"
                + str(args.latent)
                + "/bond"
                + str(args.bond_dim)
                + "/run_"
                + str(args.run)
            )
    else:
        save_dir = (
            args.save_dir
            + "/"
            + args.initializer
            + "/10k/"
            + str(args.embedding)
            + "/lat"
            + str(args.latent)
            + "/bond"
            + str(args.bond_dim)
            + "/run_"
            + str(args.run)
        )

    # set standardization and minmax to bool
    standardization = args.standardization == "yes"

    minmax = args.minmax == "yes"

    # check result dir
    if not os.path.exists(save_dir):  # noqa: PTH110
        # Create a new directory because it does not exist
        os.makedirs(save_dir)  # noqa: PTH103

    if args.seed is not None:
        # Use specified seed for reproducibility
        seed = args.seed
        print(
            Colors.YELLOW.value
            + f"Using specified seed: {seed}"
            + Colors.RESET.value
            + "\n",
            end="",
        )
    else:
        # Generate random seed for exploration
        seed = int.from_bytes(os.urandom(4), "big")
        print(
            Colors.YELLOW.value
            + f"Using random seed: {seed}"
            + Colors.RESET.value
            + "\n",
            end="",
        )

    # Set random seed
    np.random.seed(seed)  # noqa: NPY002
    key = jax.random.PRNGKey(seed)

    # Set JAX to use 64-bit precision
    # This is important for numerical stability in some cases
    jax.config.update("jax_enable_x64", True)

    # Parse embedding string to get type and degree
    embedding_string = args.embedding
    try:
        embedding_type, degree_str = embedding_string.split("_", 1)
        degree = int(degree_str)
    except ValueError:
        raise ValueError(
            Colors.RED.value
            + f"Invalid embedding format: {embedding_string}. Expected format: 'name_degree' (e.g., 'fourier_2')"
            + Colors.RESET.value
        ) from None

    # Initialize embedding based on type and degree
    if embedding_type == "fourier":
        phys_dim = (
            degree * 2
        )  # Each frequency component adds 2 dimensions (sin and cos)
        embedding: Embedding = FourierEmbedding(p=degree)
    elif embedding_type == "legendre":
        phys_dim = degree + 1  # Legendre polynomials from degree 0 to degree
        embedding = LegendreEmbedding(degree=degree)
    elif embedding_type == "laguerre":
        phys_dim = degree + 1  # Laguerre polynomials from degree 0 to degree
        embedding = LaguerreEmbedding(degree=degree)
    elif embedding_type == "hermite":
        phys_dim = degree + 1  # Hermite polynomials from degree 0 to degree
        embedding = HermiteEmbedding(degree=degree)
    else:
        raise ValueError(
            Colors.RED.value
            + f"Invalid embedding type: {embedding_type}. Supported types: fourier, legendre, laguerre, hermite"
            + Colors.RESET.value
        )

    # Set the standard deviation for the initializer
    # This is a heuristic value based on the bond dimension and physical dimension from the paper https://arxiv.org/abs/2310.20498
    std = np.power(float(phys_dim * args.bond_dim), -1)

    # Define the possible initializers
    initializers = {
        "gramschmidt_n_std": gramschmidt("normal", std, dtype=jnp.float64),
        "randn_std": tn4ml.initializers.randn(std),
        "randn_1e-2": tn4ml.initializers.randn(1e-2),
        "unitary": rand_unitary(),
    }

    # Check if the initializer is valid
    if args.initializer not in initializers:
        raise ValueError(
            Colors.RED.value
            + f"Invalid initializer: {args.initializer}. Supported initializers: {', '.join(initializers.keys())}"
            + Colors.RESET.value
        )

    print(
        Colors.BLUE.value + "Loading the model..." + Colors.RESET.value + "\n", end=""
    )
    # Save model
    model_name = "model"
    model = load_model(model_name, save_dir)

    # EVALUATION
    print(Colors.BLUE.value + "Evaluating model..." + Colors.RESET.value + "\n", end="")

    # Load scalers
    prefix = "train_qcd"
    if args.standardization == "yes":
        scaler_path = os.path.join(save_dir, f"scaler_standard_{prefix}.pkl")  # noqa: PTH118
        if os.path.exists(scaler_path):  # noqa: PTH110
            with open(scaler_path, "rb") as f:  # noqa: PTH123
                scaler = joblib.load(f)
        else:
            raise FileNotFoundError(
                Colors.RED.value
                + f"Scaler file not found: {scaler_path}"
                + Colors.RESET.value
            )
    else:
        scaler = None

    if args.minmax == "yes":
        scaler_path = os.path.join(save_dir, f"scaler_minmax_{prefix}.pkl")  # noqa: PTH118
        if os.path.exists(scaler_path):  # noqa: PTH110
            with open(scaler_path, "rb") as f:  # noqa: PTH123
                min_max_scaler = joblib.load(f)
        else:
            raise FileNotFoundError(
                Colors.RED.value
                + f"Scaler file not found: {scaler_path}"
                + Colors.RESET.value
            )
    else:
        min_max_scaler = None

    # Load test data
    read_data_dir = f"{args.load_dir}/latent{args.latent}"
    qcd_test_scaled = load_test_data(
        f"{read_data_dir}",
        dataset_type="qcd",
        scaler=scaler,
        min_max_scaler=min_max_scaler,
        test_size=args.test_size,
        shuffle_seed=seed,
    )

    sig_test_scaled = load_test_data(
        f"{read_data_dir}/latentrep_{args.signal_name}.h5",
        dataset_type="signal",
        scaler=scaler,
        min_max_scaler=min_max_scaler,
        test_size=args.test_size,
        shuffle_seed=seed,
    )

    # Calculate Fidelity - AD score
    print(
        Colors.YELLOW.value
        + "Calculating fidelity scores..."
        + Colors.RESET.value
        + "\n",
        end="",
    )

    fid_qcd = calc_fidelity_batch(
        qcd_test_scaled, model, embedding=embedding, batch_size=args.batch_size
    )

    fid_sig = calc_fidelity_batch(
        sig_test_scaled, model, embedding=embedding, batch_size=args.batch_size
    )

    # Save anomaly scores
    print(
        Colors.BLUE.value + "Saving fidelity scores..." + Colors.RESET.value + "\n",
        end="",
    )
    with h5py.File(f"{save_dir}/fidelity_scores_{args.signal_name}.h5", "w") as file:
        file.create_dataset("loss_qcd", data=fid_qcd)
        file.create_dataset("loss_sig", data=fid_sig)

pipeline.py: * View in browser: pipeline.py * View on GitHub

pipeline.py#

import argparse
import json
import os

import h5py
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
from utils import (
    Colors,
    _ensure_data_exists,
    calc_fidelity_batch,
    load_test_data,
    load_train_data,
)

import tn4ml
from tn4ml.embeddings import (
    Embedding,
    FourierEmbedding,
    HermiteEmbedding,
    LaguerreEmbedding,
    LegendreEmbedding,
)
from tn4ml.initializers import gramschmidt, rand_unitary
from tn4ml.metrics import NegLogLikelihood
from tn4ml.models.mps import MPS_initialize
from tn4ml.util import EarlyStopping, TrainingType

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="read arguments for training of TN model"
    )
    parser.add_argument(
        "-save_dir",
        dest="save_dir",
        type=str,
        help="path to directory for saving results",
        default="results/",
    )
    parser.add_argument(
        "-load_dir",
        dest="load_dir",
        type=str,
        help="path to directory for loading the data",
    )

    # data params
    parser.add_argument(
        "-feature_range",
        dest="feature_range",
        type=float,
        nargs=2,
        default=[0, 1],
        help="Feature range for scaling",
    )
    parser.add_argument(
        "-seed", dest="seed", type=int, help="Seed for random number generator"
    )
    parser.add_argument(
        "-standardization",
        dest="standardization",
        type=str,
        default="yes",
        choices=["yes", "no"],
        help="Standardization of data",
    )
    parser.add_argument(
        "-minmax",
        dest="minmax",
        type=str,
        default="yes",
        choices=["yes", "no"],
        help="Minmax scaling of data",
    )
    parser.add_argument(
        "-embedding", dest="embedding", type=str, help="Embedding type for input data"
    )
    parser.add_argument("-test_size", dest="test_size", type=int, help="Test size")
    parser.add_argument("-train_size", dest="train_size", type=int, help="Train size")
    parser.add_argument(
        "-signal_name",
        dest="signal_name",
        type=str,
        default="RSGraviton_WW_NA_35",
        help="Name of signal",
    )

    # MPS params
    parser.add_argument(
        "-bond_dim", dest="bond_dim", type=int, default=5, help="Bond dimension"
    )
    parser.add_argument(
        "-initializer", dest="initializer", type=str, help="Type of MPS initialization"
    )
    parser.add_argument("-shape_method", dest="shape_method", type=str, default="even")

    # training params
    parser.add_argument("-lr", dest="lr", type=float, default=1e-3)
    parser.add_argument("-min_delta", dest="min_delta", type=float, default=0)
    parser.add_argument("-patience", dest="patience", type=int, default=20)
    parser.add_argument("-epochs", dest="epochs", type=int, default=100)
    parser.add_argument("-batch_size", dest="batch_size", type=int, default=32)
    parser.add_argument(
        "-run", dest="run", type=int, help="Number of training repetitions"
    )

    parser.add_argument(
        "-latent", dest="latent", type=int, help="Latent space dimension"
    )

    args = parser.parse_args()
    params = vars(args)

    # Get paths to all data files, downloading them if necessary
    print(
        Colors.YELLOW.value + "Checking data folder..." + Colors.RESET.value + "\n",
        end="",
    )

    _ensure_data_exists(args.load_dir, args.latent)

    print(Colors.BLUE.value + "Importing data... " + Colors.RESET.value + "\n", end="")

    if args.standardization == "yes":
        save_dir = (
            args.save_dir
            + "/"
            + args.initializer
            + "/10k_standard/"
            + str(args.embedding)
            + "/lat"
            + str(args.latent)
            + "/bond"
            + str(args.bond_dim)
            + "/run_"
            + str(args.run)
        )
    elif args.minmax == "yes":
        if tuple(args.feature_range) == (-1, 1):
            save_dir = (
                args.save_dir
                + "/"
                + args.initializer
                + "/10k_minmax-11/"
                + str(args.embedding)
                + "/lat"
                + str(args.latent)
                + "/bond"
                + str(args.bond_dim)
                + "/run_"
                + str(args.run)
            )
        else:
            save_dir = (
                args.save_dir
                + "/"
                + args.initializer
                + "/10k_minmax01/"
                + str(args.embedding)
                + "/lat"
                + str(args.latent)
                + "/bond"
                + str(args.bond_dim)
                + "/run_"
                + str(args.run)
            )
    else:
        save_dir = (
            args.save_dir
            + "/"
            + args.initializer
            + "/10k/"
            + str(args.embedding)
            + "/lat"
            + str(args.latent)
            + "/bond"
            + str(args.bond_dim)
            + "/run_"
            + str(args.run)
        )

    # set standardization and minmax to bool
    standardization = args.standardization == "yes"

    minmax = args.minmax == "yes"

    # check result dir
    if not os.path.exists(save_dir):  # noqa: PTH110
        # Create a new directory because it does not exist
        os.makedirs(save_dir)  # noqa: PTH103

    if args.seed is not None:
        # Use specified seed for reproducibility
        seed = args.seed
        print(
            Colors.BLUE.value
            + f"Using specified seed: {seed}"
            + Colors.RESET.value
            + "\n",
            end="",
        )
    else:
        # Generate random seed for exploration
        seed = int.from_bytes(os.urandom(4), "big")
        print(
            Colors.BLUE.value
            + f"Using random seed: {seed}"
            + Colors.RESET.value
            + "\n",
            end="",
        )

    # Set random seed
    np.random.seed(seed)  # noqa: NPY002
    key = jax.random.PRNGKey(seed)

    # Set JAX to use 64-bit precision
    # This is important for numerical stability in some cases
    jax.config.update("jax_enable_x64", True)

    # Load_data
    read_data_file = f"{args.load_dir}/latent{args.latent}/latentrep_QCD_sig.h5"
    train_data, scalers = load_train_data(
        read_data_file,
        args.train_size,
        minmax,
        standardization,
        feature_range=tuple(args.feature_range),
        shuffle_seed=seed,
        save_dir=save_dir,
    )

    # Create a MPS model
    L = train_data.shape[1]
    print(
        Colors.BLUE.value + "Number of tensors: " + str(L) + Colors.RESET.value + "\n",
        end="",
    )

    # Parse embedding string to get type and degree
    embedding_string = args.embedding
    try:
        embedding_type, degree_str = embedding_string.split("_", 1)
        degree = int(degree_str)
    except ValueError:
        raise ValueError(  # noqa: B904
            f"Invalid embedding format: {embedding_string}. Expected format: 'name_degree' (e.g., 'fourier_2')"
        )

    # Initialize embedding based on type and degree
    if embedding_type == "fourier":
        phys_dim = (
            degree * 2
        )  # Each frequency component adds 2 dimensions (sin and cos)
        embedding: Embedding = FourierEmbedding(p=degree)
    elif embedding_type == "legendre":
        phys_dim = degree + 1  # Legendre polynomials from degree 0 to degree
        embedding = LegendreEmbedding(degree=degree)
    elif embedding_type == "laguerre":
        phys_dim = degree + 1  # Laguerre polynomials from degree 0 to degree
        embedding = LaguerreEmbedding(degree=degree)
    elif embedding_type == "hermite":
        phys_dim = degree + 1  # Hermite polynomials from degree 0 to degree
        embedding = HermiteEmbedding(degree=degree)
    else:
        raise ValueError(
            f"Invalid embedding type: {embedding_type}. Supported types: fourier, legendre, laguerre, hermite"
        )

    print(
        Colors.BLUE.value
        + f"Using {embedding_type} embedding with degree {degree} (physical dimension: {phys_dim})"
        + Colors.RESET.value
        + "\n",
        end="",
    )

    # Set the standard deviation for the initializer
    # This is a heuristic value based on the bond dimension and physical dimension from the paper https://arxiv.org/abs/2310.20498
    std = np.power(float(phys_dim * args.bond_dim), -1)

    # Define the possible initializers
    initializers = {
        "gramschmidt_n_std": gramschmidt("normal", std, dtype=jnp.float64),
        "randn_std": tn4ml.initializers.randn(std),
        "randn_1e-2": tn4ml.initializers.randn(1e-2),
        "unitary": rand_unitary(),
    }

    # Check if the initializer is valid
    if args.initializer not in initializers:
        raise ValueError(
            f"Invalid initializer: {args.initializer}. Supported initializers: {', '.join(initializers.keys())}"
        )

    # Initialize the MPS model
    shape_method = (
        args.shape_method
    )  # shape method defines how the tensors are arranged in the MPS
    compress = (
        False  # compress the MPS tensors - not used in this example, feature in quimb
    )
    add_identity = False  # option to add identity tensors to the MPS
    canonical_center = 0  # canonical center at the first tensor
    canonize = (True, 0)  # flag to canonize the MPS tensors during training

    print(
        Colors.BLUE.value + "Initializing MPS model..." + Colors.RESET.value + "\n",
        end="",
    )
    model = MPS_initialize(
        L=L,
        initializer=initializers[args.initializer],
        key=key,
        shape_method=shape_method,
        bond_dim=args.bond_dim,
        phys_dim=phys_dim,
        cyclic=False,
        compress=compress,
        add_identity=add_identity,
        canonical_center=canonical_center,
        boundary="obc",
    )

    # Define training parameters
    optimizer = optax.adam
    strategy = "global"
    loss = NegLogLikelihood
    train_type = TrainingType.UNSUPERVISED
    learning_rate = args.lr

    # Configure the model with the optimizer, strategy, loss function, and training type
    model.configure(
        optimizer=optimizer,
        strategy=strategy,
        loss=loss,
        train_type=train_type,
        learning_rate=learning_rate,
    )

    # Initialize the early stopping callback
    earlystop = EarlyStopping(
        min_delta=args.min_delta, patience=args.patience, mode="min", monitor="loss"
    )

    # Train the model
    print(Colors.BLUE.value + "Training model..." + Colors.RESET.value + "\n", end="")
    history = model.train(
        train_data,
        epochs=args.epochs,
        batch_size=args.batch_size,
        embedding=embedding,
        normalize=True,
        dtype=jnp.float64,
        earlystop=earlystop,
        canonize=canonize,
        seed=seed,
        shuffle=True,
    )

    # -------- SAVE RESULTS AND MODEL -------------

    print(Colors.BLUE.value + "Saving the model..." + Colors.RESET.value + "\n", end="")
    # Save model
    model_name = "model"
    model.save(model_name, save_dir)

    # Plot loss
    plt.figure()
    plt.plot(range(len(history["loss"])), history["loss"], label="train")
    plt.legend()
    plt.savefig(save_dir + "/loss.pdf")

    # Save loss
    np.save(save_dir + "/loss.npy", history["loss"])

    params_save = {
        # MPS parameters
        "bond_dim": str(args.bond_dim),
        "phys_dim": str(phys_dim),
        "initializer": str(args.initializer),
        "shape_method": str(shape_method),
        "compress": str(compress),
        "add_identity": str(add_identity),
        "boundary": "obc",
        "std": str(std),
        # Data parameters
        "embedding": str(embedding_string),
        "train_size": str(args.train_size),
        "test_size": str(args.test_size),
        "signal_name": str(args.signal_name),
        "feature_range": str(args.feature_range),
        "standardization": str(args.standardization),
        "minmax": str(args.minmax),
        # Training parameters
        "learning_rate": str(args.lr),
        "batch_size": str(args.batch_size),
        "epochs": str(args.epochs),
        "patience": str(args.patience),
        "min_delta": str(args.min_delta),
        # Model configuration
        "strategy": strategy,
        "optimizer": "adam",
        "loss": "NegLogLikelihood",
        "train_type": "unsupervised",
        # Seed and paths
        "seed": str(seed),
        "save_dir": save_dir,
        "load_dir": args.load_dir,
        "latent_space_dim": str(args.latent),
    }

    # Save parameters
    print(
        Colors.BLUE.value + "Saving parameters..." + Colors.RESET.value + "\n", end=""
    )
    with open(os.path.join(save_dir, "parameters.json"), "w") as f:  # noqa: PTH118, PTH123
        json.dump(params_save, f, indent=4)

    # EVALUATION
    print(Colors.BLUE.value + "Evaluating model..." + Colors.RESET.value + "\n", end="")

    scaler = scalers["standard"] if args.standardization == "yes" else None

    min_max_scaler = scalers["minmax"] if args.minmax == "yes" else None

    # Load test data
    read_data_dir = f"{args.load_dir}/latent{args.latent}"
    qcd_test_scaled = load_test_data(
        f"{read_data_dir}",
        dataset_type="qcd",
        scaler=scaler,
        min_max_scaler=min_max_scaler,
        test_size=args.test_size,
        shuffle_seed=seed,
    )

    sig_test_scaled = load_test_data(
        f"{read_data_dir}/latentrep_{args.signal_name}.h5",
        dataset_type="signal",
        scaler=scaler,
        min_max_scaler=min_max_scaler,
        test_size=args.test_size,
        shuffle_seed=seed,
    )

    # Calculate Fidelity - AD score
    print(
        Colors.YELLOW.value
        + "Calculating fidelity scores..."
        + Colors.RESET.value
        + "\n",
        end="",
    )

    fid_qcd = calc_fidelity_batch(
        qcd_test_scaled, model, embedding=embedding, batch_size=args.batch_size
    )

    fid_sig = calc_fidelity_batch(
        sig_test_scaled, model, embedding=embedding, batch_size=args.batch_size
    )

    # Save anomaly scores
    print(
        Colors.BLUE.value + "Saving fidelity scores..." + Colors.RESET.value + "\n",
        end="",
    )
    with h5py.File(f"{save_dir}/fidelity_scores_{args.signal_name}.h5", "w") as file:
        file.create_dataset("loss_qcd", data=fid_qcd)
        file.create_dataset("loss_sig", data=fid_sig)

plot.py: * View in browser: plot.py * View on GitHub

plot.py#

import os
from collections.abc import Collection

import h5py
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.font_manager import FontProperties
from sklearn.metrics import auc
from utils import *

from tn4ml.eval import *


def load_anomaly_scores(
    signal_name: str,
    initializers_strings: Collection[str],
    latent_spaces: Collection[int],
    bond_dim: dict,
    embedding: str,
    nruns: int,
    save_dir: str,
    train_scaling: str,
):
    """
    Load anomaly scores for different initializers and bond dimensions.

    Parameters
    ----------
    signal_name : str
        Name of the signal for which anomaly scores are obtained
    initializers_strings : Collection[str]
        List of initializer names
    latent_spaces : Collection[int]
        List of latent space dimensions to compare
    bond_dim : dict
        Dictionary containing bond dimensions for each latent space - e.g. {'4': [2, 4], '8': [2, 4, 8]}
    embedding : str
        Embedding name for the model
    nruns : int
        Number of runs to average over
    save_dir : str
        Directory to save the plot
    train_scaling : str
        Training size and scaling to consider

    Returns
    -------
    tpr_per_init : dict
        Dictionary containing true positive rates for each initializer and bond dimension
    tpr_per_init_err : dict
        Dictionary containing statistical errors for true positive rates - from multiple runs
    fpr_per_init : dict
        Dictionary containing false positive rates for each initializer and bond dimension
    fpr_per_init_err : dict
        Dictionary containing statistical errors for false positive rates - from multiple runs
    auc_per_init : dict
        Dictionary containing area under the curve values for each initializer and bond dimension
    auc_per_init_err : dict
        Dictionary containing statistical errors for area under the curve values - from multiple runs
    fpr_per_tpr_8_per_init : dict
        Dictionary containing false positive rates for TPR = 0.8 for each initializer and bond dimension
    fpr_per_tpr_8_per_init_err : dict
        Dictionary containing statistical errors for false positive rates for TPR = 0.8 - from multiple runs
    fpr_per_tpr_6_per_init : dict
        Dictionary containing false positive rates for TPR = 0.6 for each initializer and bond dimension
    fpr_per_tpr_6_per_init_err : dict
        Dictionary containing statistical errors for false positive rates for TPR = 0.6 - from multiple runs
    """
    # Initialize dictionaries to store results
    tpr_per_init = {}
    tpr_per_init_err = {}
    fpr_per_init = {}
    fpr_per_init_err = {}
    auc_per_init = {}
    auc_per_init_err = {}
    fpr_per_tpr_8_per_init = {}
    fpr_per_tpr_8_per_init_err = {}
    fpr_per_tpr_6_per_init = {}
    fpr_per_tpr_6_per_init_err = {}

    for initializer in initializers_strings:
        for _i, lat in enumerate(latent_spaces):
            for _j, bond in enumerate(bond_dim[str(lat)]):
                loss_qcd_runs = []
                loss_sig_runs = []
                tpr_data = []
                fpr_data = []
                auc_data = []
                fpr_per_tpr_8_data = []
                fpr_per_tpr_6_data = []
                for run in range(1, nruns + 1):
                    path = f"{save_dir}/{initializer}/{train_scaling}/{embedding}/lat{lat}/bond{bond}/run_{run}/fidelity_scores_{signal_name}.h5"
                    if not os.path.exists(path):  # noqa: PTH110
                        continue
                    with h5py.File(path, "r") as file:
                        loss_qcd = file["loss_qcd"][:]
                        loss_sig = file["loss_sig"][:]

                        loss_qcd = np.power(loss_qcd, 2)
                        loss_sig = np.power(loss_sig, 2)

                        loss_qcd_runs.append(loss_qcd)
                        loss_sig_runs.append(loss_sig)

                        fpr, tpr = get_roc_curve_data(
                            loss_sig, loss_qcd, anomaly_det=True
                        )
                        tpr_data.append(tpr)
                        fpr_data.append(fpr)
                        # Get auc
                        auc_value = auc(fpr, tpr)
                        auc_data.append(auc_value)

                        # Get fpr per tpr = {0.8, 0.6}
                        fpr_per_tpr_8 = get_FPR_for_fixed_TPR(
                            0.8, np.array(fpr), np.array(tpr), tolerance=0.01
                        )
                        fpr_per_tpr_6 = get_FPR_for_fixed_TPR(
                            0.6, np.array(fpr), np.array(tpr), tolerance=0.01
                        )
                        fpr_per_tpr_8_data.append(fpr_per_tpr_8)
                        fpr_per_tpr_6_data.append(fpr_per_tpr_6)

                loss_qcd = get_mean_and_error(np.array(loss_qcd_runs))
                loss_sig = get_mean_and_error(np.array(loss_sig_runs))

                if np.isnan(loss_qcd[0]).sum() > 0:
                    print(f"{Colors.RED.value}{path}: NaNs{Colors.RESET.value}")
                    continue

                # Get mean error for tpr, fpr
                tpr_mean_error = get_mean_and_error(np.array(tpr_data))
                tpr_per_init[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = tpr_mean_error[0]
                tpr_per_init_err[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = tpr_mean_error[1]

                fpr_mean_error = get_mean_and_error(1.0 / np.array(fpr_data))
                fpr_per_init[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = fpr_mean_error[0]
                fpr_per_init_err[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = fpr_mean_error[1]

                # AUC mean error
                auc_mean_error = get_mean_and_error(np.array(auc_data))
                auc_per_init[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = auc_mean_error[0]
                auc_per_init_err[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = auc_mean_error[1]

                # fpr per tpr = 0.8
                fpr_per_tpr_8_mean_error = get_mean_and_error(
                    1.0 / np.array(fpr_per_tpr_8_data)
                )
                fpr_per_tpr_8_per_init[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = fpr_per_tpr_8_mean_error[0]
                fpr_per_tpr_8_per_init_err[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = fpr_per_tpr_8_mean_error[1]

                # fpr per tpr = 0.6
                fpr_per_tpr_6_mean_error = get_mean_and_error(
                    1.0 / np.array(fpr_per_tpr_6_data)
                )
                fpr_per_tpr_6_per_init[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = fpr_per_tpr_6_mean_error[0]
                fpr_per_tpr_6_per_init_err[
                    f"init={initializer},bond={bond},lat={lat},s={signal_name}"
                ] = fpr_per_tpr_6_mean_error[1]

    return (
        tpr_per_init,
        tpr_per_init_err,
        fpr_per_init,
        fpr_per_init_err,
        auc_per_init,
        auc_per_init_err,
        fpr_per_tpr_8_per_init,
        fpr_per_tpr_8_per_init_err,
        fpr_per_tpr_6_per_init,
        fpr_per_tpr_6_per_init_err,
    )


def plot_losses_per_initializer(
    latent: int,
    bond_dims: Collection[int],
    initializers_strings: Collection[str],
    embedding: str,
    save_dir: str,
    N_epochs: int = 1000,
    train_size: str = "10k",
    minmax: str = "minmax-11",
    nruns: int = 5,
):
    """Create a subplot grid with training loss plots for all initializers for fixed embedding.

    Parameters
    ----------
    latent : int
        Latent space dimension
    bond_dims : Collection[int]
        List of bond dimensions to compare
    initializers_strings : Collection[str]
        List of initializer names to compare
    embedding : str
        Embedding name for the model
    save_dir : str
        Directory to save the plot
    N_epochs : int, optional
        Number of epochs for training, default 1000
    train_size : str, optional
        Training size to consider, default "10k"
    minmax : str, optional
        Min-Max scaling option, default "minmax-11"
    nruns : int, optional
        Number of runs to average over, default 5

    Returns
    -------
    None

    Notes
    -----
    - The function creates a grid of subplots, each showing the training loss for a different initializer.
    - Each subplot contains multiple lines representing different bond dimensions - corresponding to different colors in a color palette.
    - The function is assuming the data files are structured in a specific way, but you can change it to fit your needs.


    Example
    -------
    >>> plot_losses_per_initializer(
    ...     latent = 4,
    ...     bond_dims = [2, 4, 8, 16],
    ...     initializers_strings = ['unitary_canonize', 'randn_std', 'randn_1e-2', 'gramschmidt_n_std'],
    ...     embedding = 'laguerre_2',
    ...     save_dir = './results',
    ...     N = 1000,
    ...     train_size = '10k',
    ...     minmax = 'minmax-11',
    ...     nruns = 10
    ... )

    """
    # Calculate grid dimensions
    n_initializers = len(initializers_strings)
    n_cols = min(2, n_initializers)  # Maximum 4 columns
    n_rows = (n_initializers + n_cols - 1) // n_cols  # Ceiling division

    # Color maps for different bond dimensions
    color_maps = [
        "Blues",
        "Reds",
        "Greens",
        "Purples",
        "Oranges",
        "YlOrBr",
        "GnBu",
        "PuRd",
    ]

    # Create figure with subplots
    _fig, axes = plt.subplots(n_rows, n_cols, figsize=(6 * n_cols, 5 * n_rows))

    # Flatten axes array for easy indexing if multiple rows/columns
    if n_rows > 1 or n_cols > 1:  # noqa: SIM108
        axes = axes.flatten()
    else:
        axes = [axes]  # Convert to list for single subplot case

    # Plot for each initializer
    for i, initializer in enumerate(initializers_strings):
        ax = axes[i]

        # Plot for each bond dimension
        for j, bond in enumerate(bond_dims):
            # Create color spectrum for this bond dimension
            cmap = plt.get_cmap(color_maps[j % len(color_maps)])

            # Process each run
            for run in range(1, nruns + 1):
                # Load loss data
                file_path = f"{save_dir}/{initializer}/{train_size}_{minmax}/{embedding}/lat{latent}/bond{bond}/run_{run}/loss.npy"

                try:
                    # Load and pad with NaN
                    loss_data = np.load(file_path)
                    loss = np.full(N_epochs, np.nan)
                    loss[: len(loss_data)] = loss_data

                    # Plot with color from the spectrum
                    color_intensity = 0.3 + 0.6 * (run / nruns)
                    color = cmap(color_intensity)

                    # Plot data
                    epochs = np.arange(1, N_epochs + 1)
                    ax.plot(epochs, loss, "-", color=color, linewidth=3, alpha=0.7)

                except FileNotFoundError:
                    print(
                        f"{Colors.RED.value}File not found: {file_path}{Colors.RESET.value}"
                    )
                    continue

            # Add label for this bond dimension (just once per bond)
            ax.plot([], [], "-", color=cmap(0.5), label=rf"$\chi$ = {bond}")

        # Style this subplot
        ax.grid(True, linestyle="--", alpha=0.4)
        ax.set_xlabel("Epochs", fontsize=15)
        ax.set_ylabel("Training Loss", fontsize=15)

        # Create a label for the initializer
        if initializer == "gramschmidt_n_std":
            initializer = "Gram-Schmidt (normal - std=0.1667)"
        elif initializer == "randn_std":
            initializer = "Random (normal - std=0.1667)"
        elif initializer == "randn_1e-2":
            initializer = "Random (normal - std=0.01)"
        elif initializer == "unitary":
            initializer = "Unitary"

        ax.set_title(f"{initializer}", fontsize=17)
        ax.set_xlim(0, N_epochs)

        # set font size for ticks
        ax.tick_params(axis="both", which="major", labelsize=12)
        ax.tick_params(axis="both", which="minor", labelsize=12)

    axes[0].legend(loc="upper right", frameon=True, fontsize=14)
    # Hide any unused subplots
    for j in range(i + 1, n_rows * n_cols):
        if j < len(axes):
            axes[j].axis("off")

    # Hide y-axis labels for all but the first column
    for ax in axes:
        ax.label_outer()

    plt.tight_layout()

    # Save figure
    os.makedirs(  # noqa: PTH103
        f"{save_dir}/plots/{train_size}_{minmax}/{embedding}/lat{latent}/",
        exist_ok=True,
    )

    plt.savefig(
        f"{save_dir}/plots/{train_size}_{minmax}/{embedding}/lat{latent}/train_loss.pdf"
    )


def compare_anomaly_scores_per_embedding(
    embeddings: Collection[str],
    initializer: str,
    latent: int,
    bond_dims: Collection[int],
    signal_name: str,
    save_dir: str,
    train_size_scaling: Collection[str],
    nruns: int = 10,
):
    """
    Create a single plot with subplots for different embeddings, showing QCD and BSM
    anomaly scores distributions for each bond dimension with consistent coloring.

    Parameters
    ----------
    embeddings : Collection[str]
        List of embedding names to compare
    initializer : str
        Initializer name for the model
    latent : int
        Latent space dimension
    bond_dims : Collection[int]
        List of bond dimensions to compare
    signal_name : str
        Name of the signal to be used in the plot
    save_dir : str
        Directory to save the plot
    train_size_scaling : Collection[str]
        List of training sizes and scaling to consider for each embedding
    nruns : int, optional
        Number of runs to average over, default 10

    Returns
    -------
    fig : matplotlib.figure.Figure
        The figure object containing the plot
    axs : List[matplotlib.axes.Axes]
        List of axes objects for each subplot

    Notes
    -----
    - The function assumes that the data files are structured in a specific way,
        with paths containing the bond dimension and run number.
    - The function uses a consistent color scheme for different bond dimensions across all subplots.
    - The function handles multiple runs by averaging the results and plotting the distributions.


    Example
    -------
    >>> compare_anomaly_scores_per_embedding(
    ...     embeddings = ['laguerre_2', 'legendre_2', 'hermite_2'],
    ...     initializer = 'unitary',
    ...     latent = 4,
    ...     bond_dims = [2, 4, 8, 16],
    ...     signal_name = 'AtoHZ_to_ZZZ_35',
    ...     save_dir = './results',
    ...     train_size_scaling = ['10k_minmax01', '10k_minmax-11', '10k_minmax-11'],
    ...     nruns = 10
    ... )
    """  # noqa: D205
    # Set up color maps for different bond dimensions - using brighter colors
    color_maps = [
        "Blues",
        "Reds",
        "Greens",
        "Purples",
        "Oranges",
        "YlOrBr",
        "GnBu",
        "PuRd",
    ]

    # Create figure with subplots (one per embedding)
    fig, axs = plt.subplots(
        1, len(embeddings), figsize=(6 * len(embeddings), 5), sharey=True
    )

    # Handle case of single embedding
    if len(embeddings) == 1:
        axs = [axs]

    # Create a color dictionary with brighter colors (using 0.8 instead of 0.6)
    bond_colors = {
        bond: plt.get_cmap(color_maps[i % len(color_maps)])(0.8)
        for i, bond in enumerate(bond_dims)
    }

    # Keep track of handles for the bond dimension legend
    bond_handles = []

    # Process each embedding
    for e, embedding in enumerate(embeddings):
        train_scaling = list(train_size_scaling)[e]
        ax = axs[e]

        # Process each bond dimension
        for bond in bond_dims:
            loss_qcd_runs = []
            loss_sig_runs = []

            # Process multiple runs
            for run in range(1, nruns + 1):
                path = f"{save_dir}/{initializer}/{train_scaling}/{embedding}/lat{latent}/bond{bond}/run_{run}/fidelity_scores_{signal_name}.h5"

                if not os.path.exists(path):  # noqa: PTH110
                    continue

                # Load and process data
                with h5py.File(path, "r") as file:
                    loss_qcd = file["loss_qcd"][:]
                    loss_sig = file["loss_sig"][:]

                    # Square the losses (as in original code)
                    loss_qcd = np.power(loss_qcd, 2)
                    loss_sig = np.power(loss_sig, 2)

                    loss_qcd_runs.append(loss_qcd)
                    loss_sig_runs.append(loss_sig)

            # Skip if no valid runs found
            if not loss_qcd_runs or not loss_sig_runs:
                continue

            # Compute mean and error
            loss_qcd = get_mean_and_error(np.array(loss_qcd_runs))
            loss_sig = get_mean_and_error(np.array(loss_sig_runs))

            # Skip if NaNs detected
            if np.isnan(loss_qcd[0]).sum() > 0:
                print(
                    f"{Colors.RED.value}NaNs detected for {embedding}, bond={bond}{Colors.RESET.value}"
                )
                continue

            # Plot QCD distribution
            ax.hist(
                loss_qcd[0],
                bins=100,
                fill=False,
                color=bond_colors[bond],
                histtype="step",
                linewidth=2,
                alpha=1.0,
                density=True,
            )

            # Plot BSM distribution
            ax.hist(
                loss_sig[0],
                bins=100,
                fill=True,
                color=bond_colors[bond],
                histtype="step",
                linewidth=2,
                alpha=0.4,
                linestyle="--",
                density=True,
            )

            # Save handles for bond dimension legend (only once per bond dimension)
            if e == 0:
                bond_handles.append(
                    plt.Line2D(
                        [0],
                        [0],
                        color=bond_colors[bond],
                        linewidth=2,
                        label=f"χ = {bond}",
                    )
                )

        # Style the subplot
        if embedding == "legendre_2":
            embedding_label = r"Legendre"
        elif embedding == "fourier_2":
            embedding_label = r"Fourier"
        elif embedding == "hermite_2":
            embedding_label = r"Hermite"
        elif embedding == "laguerre_2_old_2":
            embedding_label = r"Laguerre"
        else:
            embedding_label = embedding

        ax.set_title(f"{embedding_label}", fontsize=25)
        ax.set_xlabel(r"Anomaly Score", fontsize=20)
        ax.set_yscale("log")
        ax.set_xlim(-0.01, 0.6)

        # Only add y-label to first subplot
        if e == 0:
            ax.set_ylabel("Probability Density", fontsize=20)

    # Create style legend handles that accurately represent the visualization
    style_handles = [
        mpatches.Patch(
            edgecolor="black", facecolor="none", linewidth=2, label="QCD (Background)"
        ),
        mpatches.Patch(facecolor="gray", alpha=0.4, label="BSM (Signal)"),
    ]
    # set tick size
    for ax in axs:
        ax.tick_params(axis="both", which="major", labelsize=15)
        ax.tick_params(axis="both", which="minor", labelsize=15)

    # Add bond dimension legend at the figure level
    fig.legend(
        handles=bond_handles,
        loc="upper left",
        bbox_to_anchor=(0.05, 0.87),
        frameon=True,
        fontsize=10,
        title="Bond Dimension",
    )

    # Add the QCD/BSM legend to the figure
    fig.legend(
        handles=style_handles,
        loc="upper right",
        bbox_to_anchor=(0.35, 0.87),
        frameon=True,
        fontsize=10,
        title="Distribution Type",
    )

    plt.tight_layout(
        rect=(0, 0.05, 1, 0.95)
    )  # Adjust layout to make room for the legends

    # Create save directory if it doesn't exist
    os.makedirs(f"{save_dir}/plots/comparisons/", exist_ok=True)  # noqa: PTH103

    # Save figure
    plt.savefig(
        f"{save_dir}/plots/comparisons/embedding_comparison_{initializer}_lat{latent}_{signal_name}.pdf",
        bbox_inches="tight",
    )

    return fig, axs


def compare_ROCs_per_bond(  # noqa: N802
    latent: int,
    bond_dim: Collection[int],
    initializer: str,
    initializer_string: str,  # noqa: ARG001
    tpr_per_init: dict,
    tpr_per_init_err: dict,  # noqa: ARG001
    fpr_per_init: dict,
    fpr_per_init_err: dict,
    auc_per_init: dict,
    auc_per_init_err: dict,
    save_dir: str,
    embedding: str,
    train_scaling: str,
    signal_name: str,
):
    """
    Create ROC curve plot for a single initializer with consistent coloring for bond dimensions.

    Parameters
    ----------
    latent : int
        Latent space dimension
    bond_dim : Collection[int]
        List of bond dimensions to compare
    initializer : str
        Initializer of the model
    initializer_string : str
        String representation of the initializer
    tpr_per_init : dict
        Dictionary containing true positive rates for each bond dimension
    tpr_per_init_err : dict
        Dictionary containing statistical errors for true positive rates - from multiple runs
    fpr_per_init : dict
        Dictionary containing false positive rates for each bond dimension
    fpr_per_init_err : dict
        Dictionary containing statistical errors for false positive rates - from multiple runs
    auc_per_init : dict
        Dictionary containing area under the curve values for each bond dimension
    auc_per_init_err : dict
        Dictionary containing statistical errors for area under the curve values - from multiple runs
    save_dir : str
        Directory to save the plot
    embedding : str
        Embedding name for the model
    train_scaling : str
        Training size and scaling to consider
    signal_name : str
        Name of the signal to be used in the plot

    Returns
    -------
    fig : matplotlib.figure.Figure
        The figure object containing the plot
    ax : matplotlib.axes.Axes
        The axes object for the plot
    """
    # Set up color maps for different bond dimensions - using brighter colors
    color_maps = [
        "Blues",
        "Reds",
        "Greens",
        "Purples",
        "Oranges",
        "YlOrBr",
        "GnBu",
        "PuRd",
    ]

    # Create a color dictionary with brighter colors for consistent coloring
    bond_colors = {
        bond: plt.get_cmap(color_maps[i % len(color_maps)])(0.8)
        for i, bond in enumerate(bond_dim)
    }

    # Create a single figure
    fig, ax = plt.subplots(figsize=(7, 7))

    for _j, bond in enumerate(bond_dim):
        key = f"init={initializer},bond={bond},lat={latent},s={signal_name}"

        if key not in tpr_per_init:
            print(f"{Colors.RED.value}{key} not found{Colors.RESET.value}")
            continue

        tpr = tpr_per_init[key]
        # tpr_err = tpr_per_init_err[key]  # noqa: ERA001
        fpr = fpr_per_init[key]
        fpr_err = fpr_per_init_err[key]
        auc_value = auc_per_init[key]
        auc_err = auc_per_init_err[key]

        if signal_name == "RSGraviton_WW_NA_35":  # uncertainties are bigger for G_NA
            band_ind = np.where(tpr > 0.6)[0]
        else:
            band_ind = np.where(tpr > 0.35)[0]

        # Use consistent color based on bond dimension
        ax.plot(
            tpr,
            fpr,
            label=rf"$\chi$ = {bond} ({auc_value * 100.0:.2f})$\pm$({auc_err * 100.0:.2f})",
            linewidth=2,
            color=bond_colors[bond],
        )  # Use consistent color

        # Error calculation in log space
        log_fpr = np.log10(fpr)
        rel_err = fpr_err / fpr  # Relative error
        log_err = (0.434) * rel_err  # Convert to log10 error (0.434 = 1/ln(10))

        # Calculate bounds in log space, then convert back
        log_upper = log_fpr - log_err
        log_lower = log_fpr + log_err
        fpr_upper = 10**log_upper  # convert back to linear scale
        fpr_lower = 10**log_lower  # convert back to linear scale

        # Add error bands with matching color
        ax.fill_between(
            tpr[band_ind],
            fpr_lower[band_ind],
            fpr_upper[band_ind],
            alpha=0.2,
            color=bond_colors[bond],
        )  # Match fill color

    # Add vertical dotted lines at TPR = 0.6 and TPR = 0.8
    ax.axvline(x=0.6, color="black", linestyle=":", linewidth=1.5, alpha=0.3)
    ax.axvline(x=0.8, color="black", linestyle=":", linewidth=1.5, alpha=0.3)

    # Style the plot
    # ax.set_title(f'Latent = {latent}, {initializer_string}', fontsize=16)  # noqa: ERA001
    ax.legend(loc="lower left", fontsize=12)
    ax.set_yscale("log")
    ax.set_xlabel("TPR", fontsize=20)
    ax.set_ylabel("FPR$^{-1}$", fontsize=20)
    ax.set_xticks(np.arange(0, 1.1, 0.2))
    ax.tick_params(axis="both", which="major", labelsize=14)
    ax.tick_params(axis="both", which="minor", labelsize=14)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()

    # Create directory if it doesn't exist
    os.makedirs(  # noqa: PTH103
        f"{save_dir}/plots/{train_scaling}/{embedding}/lat{latent}", exist_ok=True
    )

    # Save figure
    plt.savefig(
        f"{save_dir}/plots/{train_scaling}/{embedding}/lat{latent}/roc_curve_{signal_name}.pdf"
    )

    return fig, ax


def compare_ROC_by_signal(  # noqa: N802
    signal_names: Collection[str],
    signal_labels: Collection[str],
    latent: int,
    bond_dim: Collection[int],
    initializer: str,
    initializer_string: str,  # noqa: ARG001
    tpr_per_init: dict,
    tpr_per_init_err: dict,  # noqa: ARG001
    fpr_per_init: dict,
    fpr_per_init_err: dict,
    auc_per_init: dict,
    auc_per_init_err: dict,
    save_dir: str,
    embedding: str,
    train_scaling: str,
):
    r"""
    Compare ROC curves for different signal types with fixed model parameters.

    Parameters
    ----------
    signal_names : Collection[str]
        List of signal names to compare
    signal_labels : Collection[str]
        List of signal labels for the legend
    latent : int
        Latent space dimension
    bond_dim : Collection[int]
        List of bond dimensions to compare
    initializer : str
        Initializer of the model
    initializer_string : str
        String representation of the initializer
    tpr_per_init : dict
        Dictionary containing true positive rates for each signal
    tpr_per_init_err : dict
        Dictionary containing statistical errors for true positive rates - from multiple runs
    fpr_per_init : dict
        Dictionary containing false positive rates for each signal
    fpr_per_init_err : dict
        Dictionary containing statistical errors for false positive rates - from multiple runs
    auc_per_init : dict
        Dictionary containing area under the curve values for each signal
    auc_per_init_err : dict
        Dictionary containing statistical errors for area under the curve values - from multiple runs
    save_dir : str
        Directory to save the plot
    embedding : str
        Embedding name for the model
    train_scaling : str
        Training size and scaling to consider

    Returns
    -------
    fig : matplotlib.figure.Figure
        The figure object containing the plot
    ax : matplotlib.axes.Axes
        The axes object for the plot

    Example
    -------
    >>> compare_ROC_by_signal(
    ...     signal_names = ['RSGraviton_WW_NA_35', 'AtoHZ_to_ZZZ_35', 'RSGraviton_WW_BR_15'],
    ...     signal_labels = [r'Narrow $G \rightarrow WW$', r'$A \rightarrow HZ \rightarrow ZZZ$', r'Broad $G \rightarrow WW$'],
    ...     latent = 4,
    ...     bond_dim = [2, 4, 8, 16],
    ...     initializer = 'unitary',
    ...     initializer_string = 'Unitary',
    ...     tpr_per_init = tpr_per_init,
    ...     tpr_per_init_err = tpr_per_init_err,
    ...     fpr_per_init = fpr_per_init,
    ...     fpr_per_init_err = fpr_per_init_err,
    ...     auc_per_init = auc_per_init,
    ...     auc_per_init_err = auc_per_init_err,
    ...     save_dir = './results',
    ...     embedding = 'laguerre_2',
    ...     train_scaling = '10k_minmax-11'
    ... )

    """
    palette = ["#4CA64C", "#FF5733", "#8A2BE2"]

    fig, ax = plt.subplots(figsize=(7, 7))

    # Store handles and data for legend
    handles = []
    auc_info = []
    signal_info = []

    for i, signal_name in enumerate(signal_names):
        key = f"init={initializer},bond={bond_dim},lat={latent},s={signal_name}"

        if key not in tpr_per_init:
            print(f"{Colors.RED.value}{key} not found{Colors.RESET.value}")
            continue

        tpr = tpr_per_init[key]
        fpr = fpr_per_init[key]
        fpr_err = fpr_per_init_err[key]
        auc_value = auc_per_init[key]
        auc_err = auc_per_init_err[key]

        # Determine where to show error bands
        band_ind = np.where(tpr > 0.35)[0]
        if "RSGraviton_WW_NA" in signal_name:
            band_ind = np.where(tpr > 0.6)[0]

        # Plot the ROC curve
        (line,) = ax.plot(tpr, fpr, linewidth=2, color=palette[i])

        # Store data for legend
        handles.append(line)
        auc_info.append(f"{auc_value * 100:.2f}±{auc_err * 100:.2f}")
        signal_info.append(list(signal_labels)[i])

        # Error calculation and bands
        log_fpr = np.log10(fpr)
        rel_err = fpr_err / fpr
        log_err = 0.434 * rel_err

        log_upper = log_fpr - log_err
        log_lower = log_fpr + log_err
        fpr_upper = 10**log_upper
        fpr_lower = 10**log_lower

        ax.fill_between(
            tpr[band_ind],
            fpr_lower[band_ind],
            fpr_upper[band_ind],
            alpha=0.2,
            color=palette[i],
        )

    # Create legend with AUC values first
    legend_labels = []
    for auc_label, signal in zip(auc_info, signal_info, strict=False):
        legend_labels.append(f"{auc_label}   {signal}")

    # Add legend with AUC first
    legend = ax.legend(
        handles, legend_labels, loc="lower left", frameon=True, fontsize=12
    )

    # Add the column headers with AUC first
    legend.set_title("    AUC               BSM Scenario", prop=FontProperties(size=12))

    # Add vertical dotted lines at TPR = 0.6 and TPR = 0.8
    ax.axvline(x=0.6, color="black", linestyle=":", linewidth=1.5, alpha=0.3)
    ax.axvline(x=0.8, color="black", linestyle=":", linewidth=1.5, alpha=0.3)

    # Style the plot
    ax.set_yscale("log")
    ax.set_xlabel("TPR", fontsize=20)
    ax.set_ylabel("FPR$^{-1}$", fontsize=20)
    ax.set_xticks(np.arange(0, 1.1, 0.2))
    ax.tick_params(axis="both", which="major", labelsize=14)
    ax.tick_params(axis="both", which="minor", labelsize=14)

    # ax.set_yticks(fontsize=14)  # noqa: ERA001
    ax.grid(True, alpha=0.3)

    plt.tight_layout()

    # Create directory and save
    os.makedirs(  # noqa: PTH103
        f"{save_dir}/plots/{train_scaling}/{embedding}/lat{latent}", exist_ok=True
    )
    plt.savefig(
        f"{save_dir}/plots/{train_scaling}/{embedding}/lat{latent}/roc_curve_compare_signals.pdf"
    )

    return fig, ax


def compare_ROC_by_latent(  # noqa: N802
    latent_spaces: Collection[int],
    bond_dims: dict,
    initializer: str,
    initializer_string: str,  # noqa: ARG001
    signal_name: str,
    signal_label: str,  # noqa: ARG001
    tpr_per_init: dict,
    tpr_per_init_err: dict,  # noqa: ARG001
    fpr_per_init: dict,
    fpr_per_init_err: dict,
    auc_per_init: dict,
    auc_per_init_err: dict,
    save_dir: str,
    embedding: str,
    train_scaling: str,
):
    """
    Compare ROC curves for one signal type across different latent spaces, each with a specific bond dimension.

    Parameters
    ----------
    latent_spaces : Collection[int]
        List of latent space dimensions to compare
    bond_dims : dict
        Dictionary mapping latent space dimensions to bond dimensions
    initializer : str
        Initializer of the model
    initializer_string : str
        String representation of the initializer
    signal_name : str
        Name of the signal anomaly scores are calculated for
    signal_label : str
        Label for the signal to be used in the plot
    tpr_per_init : dict
        Dictionary containing true positive rates for each latent space
    tpr_per_init_err : dict
        Dictionary containing statistical errors for true positive rates - from multiple runs
    fpr_per_init : dict
        Dictionary containing false positive rates for each latent space
    fpr_per_init_err : dict
        Dictionary containing statistical errors for false positive rates - from multiple runs
    auc_per_init : dict
        Dictionary containing area under the curve values for each latent space
    auc_per_init_err : dict
        Dictionary containing statistical errors for area under the curve values - from multiple runs
    save_dir : str
        Directory to save the plot
    embedding : str
        Embedding name for the model
    train_scaling : str
        Training size and scaling used in training

    """
    # Colors for different latent spaces
    palette = [
        "#E69F00",  # Muted orange
        "#CC6677",  # Muted red
        "#88CCEE",  # Muted blue
        "#000000",  # Black
        "#44AA99",  # Muted teal
        "#AA4499",
    ]  # Muted purple

    if len(latent_spaces) > len(palette):
        # Generate more colors if needed
        cmap = plt.get_cmap("tab10")
        palette = [cmap(i) for i in np.linspace(0, 1, len(latent_spaces))]

    fig, ax = plt.subplots(figsize=(7, 7))

    # Store handles and data for legend
    handles = []
    auc_info = []
    config_info = []

    for i, latent in enumerate(latent_spaces):
        # Get corresponding bond dimension for this latent space
        bond_dim = bond_dims[str(latent)]

        key = f"init={initializer},bond={bond_dim},lat={latent},s={signal_name}"

        if key not in tpr_per_init:
            print(f"{key} not found")
            continue

        tpr = tpr_per_init[key]
        fpr = fpr_per_init[key]
        fpr_err = fpr_per_init_err[key]
        auc_value = auc_per_init[key]
        auc_err = auc_per_init_err[key]

        if auc_value < 0.5:
            auc_value = 1 - auc_value

        # Determine where to show error bands
        band_ind = np.where(tpr > 0.35)[0]
        if "RSGraviton_WW_NA" in signal_name:
            band_ind = np.where(tpr > 0.6)[0]

        # Plot the ROC curve
        (line,) = ax.plot(tpr, fpr, linewidth=2, color=palette[i % len(palette)])

        # Store data for legend
        handles.append(line)
        auc_info.append(f"{auc_value * 100:.2f}±{auc_err * 100:.2f}")
        config_info.append(f"lat = {latent}, χ = {bond_dim}")

        # Error calculation and bands
        log_fpr = np.log10(fpr)
        rel_err = fpr_err / fpr
        log_err = 0.434 * rel_err

        log_upper = log_fpr - log_err
        log_lower = log_fpr + log_err
        fpr_upper = 10**log_upper
        fpr_lower = 10**log_lower

        ax.fill_between(
            tpr[band_ind],
            fpr_lower[band_ind],
            fpr_upper[band_ind],
            alpha=0.2,
            color=palette[i % len(palette)],
        )

    # Create legend with AUC values first
    legend_labels = []
    for auc_label, config in zip(auc_info, config_info, strict=False):
        legend_labels.append(f"{auc_label}   {config}")

    # Add legend with AUC first
    legend = ax.legend(
        handles, legend_labels, loc="lower left", frameon=True, fontsize=12
    )

    # Add the column headers with AUC first
    legend.set_title(
        "    AUC               Configuration", prop=FontProperties(size=12)
    )

    # Add vertical dotted lines at TPR = 0.6 and TPR = 0.8
    ax.axvline(x=0.6, color="black", linestyle=":", linewidth=1.5, alpha=0.3)
    ax.axvline(x=0.8, color="black", linestyle=":", linewidth=1.5, alpha=0.3)

    # Style the plot
    ax.set_yscale("log")
    ax.set_xlabel("TPR", fontsize=20)
    ax.set_ylabel("FPR$^{-1}$", fontsize=20)
    # ax.set_title(f'Signal: {signal_label}, {initializer_string}', fontsize=16)  # noqa: ERA001
    ax.set_xticks(np.arange(0, 1.1, 0.2))
    ax.tick_params(axis="both", which="major", labelsize=14)
    ax.tick_params(axis="both", which="minor", labelsize=14)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()

    # Create directory and save
    os.makedirs(f"{save_dir}/plots/{train_scaling}/{embedding}", exist_ok=True)  # noqa: PTH103
    plt.savefig(
        f"{save_dir}/plots/{train_scaling}/{embedding}/roc_curve_latent_comparison_{signal_name}.pdf"
    )

    return fig, ax

utils.py: * View in browser: utils.py * View on GitHub

utils.py#

import os
import tarfile
from enum import Enum
from pathlib import Path

import h5py
import jax
import joblib
import numpy as np
import wget
from sklearn.preprocessing import MinMaxScaler, StandardScaler

from tn4ml.embeddings import Embedding, TrigonometricEmbedding, embed


class Colors(Enum):
    """ANSI color codes for terminal text styling."""

    RESET = "\033[0m"
    GREEN = "\033[32m"
    BLUE = "\033[34m"
    ORANGE = "\033[38;2;255;165;0m"
    PINK = "\033[38;2;255;105;180m"
    RED = "\033[31m"
    YELLOW = "\033[33m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


def _download_data(data_url: str, data_dir: str | Path = "."):
    """
    Downloads the jet data if it does not already exist.

    Parameters
    ----------
    data_url : str
        URL to the data file
    data_dir : str, optional
        Directory to save the data, default '.'

    Returns
    -------
    None
    """  # noqa: D401
    data_path = Path(data_dir)
    if not data_path.is_dir():
        os.makedirs(data_path, exist_ok=True)  # noqa: PTH103

    data_file_path = wget.download(data_url, out=str(data_path))

    data_tar = tarfile.open(data_file_path, "r:gz")  # noqa: SIM115
    data_tar.extractall(str(data_path))
    data_tar.close()
    os.remove(data_file_path)  # noqa: PTH107


def _ensure_data_exists(data_dir: str = "data", latent: int | None = None) -> None:
    """
    Check if the data directory exists and download all data if it doesn't.
    Only checks directory existence, not individual files.

    Parameters
    ----------
    data_dir : str, optional
        Base directory where data should be stored, default "data"
    latent : int, optional
        Latent space dimension used for subdirectory, default None

    Returns
    -------
    dict
        Dictionary with paths to all data files
    """  # noqa: D205
    # Create data directory path with latent dimension subfolder
    base_dir = Path(data_dir)
    data_dir_path = base_dir / f"latent{latent}" if latent is not None else base_dir

    # ONLY check if the directory exists, not individual files
    if not data_dir_path.is_dir() or not any(data_dir_path.iterdir()):
        print(
            f"{Colors.BLUE.value}Data directory {data_dir_path} does not exist. Downloading complete dataset...{Colors.RESET.value}"
            + "\n"
        )
        os.makedirs(data_dir_path, exist_ok=True)  # noqa: PTH103

        archive_url = "https://zenodo.org/records/7673769/files/QML_paper_data.tar.gz"
        try:
            _download_data(archive_url, base_dir)  # Download to base dir
            print(
                f"{Colors.BLUE.value}Archive downloaded and extracted successfully.{Colors.RESET.value}"
                + "\n"
            )
        except Exception as e:  # noqa: BLE001
            print(
                f"{Colors.RED.value}Failed to download archive: {e}{Colors.RESET.value}"
                + "\n"
            )
    else:
        print(
            f"{Colors.YELLOW.value}Data directory {data_dir_path} already exists. Assuming all data is present.{Colors.RESET.value}"
            + "\n"
        )

    return


def load_train_data(
    read_file: str,
    train_size: int = 10000,
    apply_minmax: bool = False,
    apply_standardization: bool = False,
    feature_range: tuple = (0, 1),
    shuffle_seed: int = 42,
    save_dir: str = ".",
    prefix: str = "train_qcd",
):
    """
    Load and preprocess training data from a given file.

    Parameters
    ----------
    read_file : str
        Path to the file containing the training data
    train_size : int, optional
        Number of training samples to load, default 10000
    apply_minmax : bool, optional
        Whether to apply Min-Max scaling, default False
    apply_standardization : bool, optional
        Whether to apply standardization, default False
    feature_range : tuple, optional
        The desired range for Min-Max scaling, default (0, 1)
    shuffle_seed : int, optional
        Seed for random shuffling, default 42
    save_dir : str, optional
        Directory to save scalers, default '.'
    prefix : str, optional
        Prefix for scaler filenames, default 'train_qcd'

    Returns
    -------
    np.ndarray
        Preprocessed training data
    dict
        Dictionary of fitted scalers
    """
    # Ensure the save directory exists
    os.makedirs(save_dir, exist_ok=True)  # noqa: PTH103

    # Read and prepare data
    with h5py.File(read_file, "r") as file:
        data = file["latent_space"]
        data = np.concatenate([data[:, 0, :], data[:, 1, :]], axis=-1)
        print(
            f"{Colors.BLUE.value}Input data shape: {data.shape}{Colors.RESET.value}"
            + "\n"
        )

        # Shuffle data
        np.random.seed(shuffle_seed)  # noqa: NPY002
        np.random.shuffle(data)  # noqa: NPY002

    data_train = data[:train_size]
    scalers = {}

    # Apply transformations
    if apply_standardization:
        scaler = StandardScaler()
        data_train = scaler.fit_transform(data_train)
        scalers["standard"] = scaler
        scaler_path = os.path.join(save_dir, f"scaler_standard_{prefix}.pkl")  # noqa: PTH118
        joblib.dump(scaler, scaler_path)

    if apply_minmax:
        min_max_scaler = MinMaxScaler(feature_range=feature_range)
        data_train = min_max_scaler.fit_transform(data_train)
        scalers["minmax"] = min_max_scaler
        scaler_path = os.path.join(save_dir, f"scaler_minmax_{prefix}.pkl")  # noqa: PTH118
        joblib.dump(min_max_scaler, scaler_path)

    return data_train, scalers


def load_test_data(
    read_path: str,
    dataset_type: str = "qcd",
    scaler: StandardScaler = None,
    min_max_scaler: MinMaxScaler = None,
    test_size: int = 10000,
    shuffle_seed: int = 42,
):
    """
    Load and preprocess test data from a given file.

    Parameters
    ----------
    read_path : str
        Path to the file/directory containing test data
    dataset_type : str, optional
        Type of dataset ('qcd' or 'signal'), default 'qcd'
    scaler : StandardScaler, optional
        Fitted StandardScaler to apply, default None
    min_max_scaler : MinMaxScaler, optional
        Fitted MinMaxScaler to apply, default None
    test_size : int, optional
        Number of test samples to use, default 10000
    shuffle_seed : int, optional
        Seed for random shuffling, default 42

    Returns
    -------
    np.ndarray
        Preprocessed test data
    """
    # Determine file path based on dataset type
    if dataset_type == "qcd":
        file_path = os.path.join(read_path, "latentrep_QCD_sig_testclustering.h5")  # noqa: PTH118
    else:
        file_path = read_path  # For signal, use the path directly

    # Load and prepare data
    with h5py.File(file_path, "r") as file:
        data = file["latent_space"]
        # Concatenate the two latent space components
        data_test = np.concatenate([data[:, 0, :], data[:, 1, :]], axis=-1)

        # Shuffle data
        np.random.seed(shuffle_seed)  # noqa: NPY002
        np.random.shuffle(data_test)  # noqa: NPY002
        data_test = data_test[:test_size]

    print(
        f"{Colors.BLUE.value}Input test {dataset_type} shape: {data_test.shape}{Colors.RESET.value}"
        + "\n"
    )

    # Apply transformations if provided
    if scaler is not None:
        data_test = scaler.transform(data_test)
    if min_max_scaler is not None:
        data_test = min_max_scaler.transform(data_test)

    return data_test


def calc_fidelity_batch(
    points,
    model,
    embedding: Embedding = TrigonometricEmbedding(),  # noqa: B008
    batch_size: int = 1000,
):
    """
    Calculate fidelity scores for data points in batches using vectorized operations.

    Parameters
    ----------
    points : np.ndarray
        Input data points
    model : TensorNetwork
        Trained tensor network model
    embedding : Embedding, optional
        Embedding function, default TrigonometricEmbedding()
    batch_size : int, optional
        Batch size for processing, default 1000

    Returns
    -------
    np.ndarray
        Array of fidelity scores
    """

    # Define a function that processes a single point
    def single_point_fidelity(point):
        input_mps = embed(point, embedding)
        p_mps = input_mps.H & model
        return abs(p_mps ^ all)

    # Vectorize the function for batch processing
    batch_fidelity = jax.vmap(single_point_fidelity)

    # Process in batches to avoid memory issues
    n_samples = len(points)
    n_batches = (n_samples + batch_size - 1) // batch_size
    results = []

    for i in range(n_batches):
        start_idx = i * batch_size
        end_idx = min(start_idx + batch_size, n_samples)
        batch = points[start_idx:end_idx]
        batch_results = batch_fidelity(batch)
        results.append(batch_results)

    return np.concatenate(results, axis=0)

Tensor Network for Anomaly Detection in Latent Space of Proton-Proton Collision Events at the LHC#

(optional) Download dataset#

Run training and evaluation pipeline#

Run evaluation only#

Additional Example Scripts#

This Page