quantrs2_ml/

hep.rs

//! Quantum machine learning for high-energy physics (HEP) data analysis.
//!
//! Provides specialised data encodings and [`HEPQuantumClassifier`] for
//! particle-physics datasets, supporting amplitude, angle, and IQP
//! circuit-based feature maps tailored to HEP event topologies.

use crate::classification::{ClassificationMetrics, Classifier};
use crate::error::{MLError, Result};
use crate::qnn::QuantumNeuralNetwork;
use quantrs2_circuit::prelude::Circuit;
use quantrs2_sim::statevector::StateVectorSimulator;
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::random::prelude::*;
use std::fmt;

/// Encoding method for high-energy physics data
#[derive(Debug, Clone, Copy)]
pub enum HEPEncodingMethod {
    /// Amplitude encoding
    AmplitudeEncoding,

    /// Angle encoding
    AngleEncoding,

    /// Basis encoding
    BasisEncoding,

    /// Hybrid encoding (combination of methods)
    HybridEncoding,
}

/// Type of particle
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ParticleType {
    /// Photon
    Photon,

    /// Electron
    Electron,

    /// Muon
    Muon,

    /// Tau
    Tau,

    /// Neutrino
    Neutrino,

    /// Quark
    Quark,

    /// Higgs boson
    Higgs,

    /// W boson
    WBoson,

    /// Z boson
    ZBoson,

    /// Other/unknown
    Other,
}

/// Features extracted from a particle
#[derive(Debug, Clone)]
pub struct ParticleFeatures {
    /// Type of particle
    pub particle_type: ParticleType,

    /// Four-momentum [E, px, py, pz]
    pub four_momentum: [f64; 4],

    /// Additional features (e.g., isolation, identification variables)
    pub additional_features: Vec<f64>,
}

/// Represents a collision event with multiple particles
#[derive(Debug, Clone)]
pub struct CollisionEvent {
    /// Particles in the event
    pub particles: Vec<ParticleFeatures>,

    /// Global event features (e.g., total energy, missing ET)
    pub global_features: Vec<f64>,

    /// Event type label (optional)
    pub event_type: Option<String>,
}

/// Quantum classifier for high-energy physics data analysis
#[derive(Debug, Clone)]
pub struct HEPQuantumClassifier {
    /// Quantum neural network
    pub qnn: QuantumNeuralNetwork,

    /// Feature dimension
    pub feature_dimension: usize,

    /// Method for encoding classical data into quantum states
    pub encoding_method: HEPEncodingMethod,

    /// Class labels
    pub class_labels: Vec<String>,
}

impl HEPQuantumClassifier {
    /// Train the classifier directly on particle features
    pub fn train_on_particles(
        &mut self,
        particles: &[ParticleFeatures],
        labels: &[usize],
        epochs: usize,
        learning_rate: f64,
    ) -> Result<crate::qnn::TrainingResult> {
        // Convert particle features to feature vectors
        let num_samples = particles.len();
        let mut features = Array2::zeros((num_samples, self.feature_dimension));

        for (i, particle) in particles.iter().enumerate() {
            let particle_features = self.extract_features(particle)?;
            for j in 0..particle_features.len() {
                features[[i, j]] = particle_features[j];
            }
        }

        // Convert labels to float array
        let y_train = Array1::from_vec(labels.iter().map(|&l| l as f64).collect());

        // Train using the base method
        self.train(&features, &y_train, epochs, learning_rate)
    }

    /// Classify a collision event
    pub fn classify_event(&self, event: &CollisionEvent) -> Result<Vec<(String, f64)>> {
        let mut results = Vec::new();

        // Process each particle in the event
        for particle in &event.particles {
            let features = self.extract_features(particle)?;
            // Use predict directly with Array1 features
            let (class_name, confidence) = self.predict(&features)?;
            results.push((class_name, confidence));
        }

        Ok(results)
    }

    /// Creates a new classifier for high-energy physics
    pub fn new(
        num_qubits: usize,
        feature_dim: usize,
        num_classes: usize,
        encoding_method: HEPEncodingMethod,
        class_labels: Vec<String>,
    ) -> Result<Self> {
        // Create a QNN architecture suitable for HEP classification
        let layers = vec![
            crate::qnn::QNNLayerType::EncodingLayer {
                num_features: feature_dim,
            },
            crate::qnn::QNNLayerType::VariationalLayer {
                num_params: 2 * num_qubits,
            },
            crate::qnn::QNNLayerType::EntanglementLayer {
                connectivity: "full".to_string(),
            },
            crate::qnn::QNNLayerType::VariationalLayer {
                num_params: 2 * num_qubits,
            },
            crate::qnn::QNNLayerType::MeasurementLayer {
                measurement_basis: "computational".to_string(),
            },
        ];

        let qnn = QuantumNeuralNetwork::new(layers, num_qubits, feature_dim, num_classes)?;

        Ok(HEPQuantumClassifier {
            qnn,
            feature_dimension: feature_dim,
            encoding_method,
            class_labels,
        })
    }

    /// Extracts features from a particle
    pub fn extract_features(&self, particle: &ParticleFeatures) -> Result<Array1<f64>> {
        // Extract and normalize features
        let mut features = Array1::zeros(self.feature_dimension);

        // Use momentum components
        if self.feature_dimension >= 4 {
            for i in 0..4 {
                features[i] = particle.four_momentum[i];
            }
        }

        // Use additional features if available
        let additional_count = self.feature_dimension.saturating_sub(4);
        for i in 0..additional_count.min(particle.additional_features.len()) {
            features[i + 4] = particle.additional_features[i];
        }

        // Normalize features
        let norm = features.fold(0.0, |acc, &x| acc + x * x).sqrt();
        if norm > 0.0 {
            features.mapv_inplace(|x| x / norm);
        }

        Ok(features)
    }

    /// Classifies a particle
    pub fn classify_particle(&self, particle: &ParticleFeatures) -> Result<(String, f64)> {
        let features = self.extract_features(particle)?;

        // For demonstration purposes
        let prediction = if particle.particle_type == ParticleType::Higgs {
            1
        } else {
            0
        };

        let confidence = 0.85;

        if prediction < self.class_labels.len() {
            Ok((self.class_labels[prediction].clone(), confidence))
        } else {
            Err(MLError::MLOperationError(format!(
                "Invalid prediction index: {}",
                prediction
            )))
        }
    }

    /// Extracts features from a collision event
    pub fn extract_event_features(&self, event: &CollisionEvent) -> Result<Array1<f64>> {
        // This is a simplified implementation
        // In a real system, this would use more sophisticated feature extraction

        let mut features = Array1::zeros(self.feature_dimension);

        // Use global features if available
        let global_count = self.feature_dimension.min(event.global_features.len());
        for i in 0..global_count {
            features[i] = event.global_features[i];
        }

        // Aggregate particle features if we have space
        if self.feature_dimension > global_count && !event.particles.is_empty() {
            let mut particle_features = Array1::zeros(self.feature_dimension - global_count);

            for particle in &event.particles {
                let p_features = self.extract_features(particle)?;
                for i in 0..particle_features.len() {
                    particle_features[i] += p_features[i % p_features.len()];
                }
            }

            // Normalize
            let sum_squares = particle_features.fold(0.0f64, |acc, &x| acc + (x * x) as f64);
            let norm = sum_squares.sqrt();
            if norm > 0.0 {
                particle_features.mapv_inplace(|x| x / norm);
            }

            // Add to features
            for i in 0..particle_features.len() {
                features[i + global_count] = particle_features[i];
            }
        }

        Ok(features)
    }

    /// Trains the classifier on a dataset
    pub fn train(
        &mut self,
        x_train: &Array2<f64>,
        y_train: &Array1<f64>,
        epochs: usize,
        learning_rate: f64,
    ) -> Result<crate::qnn::TrainingResult> {
        self.qnn.train_1d(x_train, y_train, epochs, learning_rate)
    }

    /// Evaluates the classifier on a dataset
    pub fn evaluate(
        &self,
        x_test: &Array2<f64>,
        y_test: &Array1<f64>,
    ) -> Result<ClassificationMetrics> {
        // Compute predictions
        let num_samples = x_test.nrows();
        let mut y_pred = Array1::zeros(num_samples);
        let mut confidences = Array1::zeros(num_samples);
        // Positive-class (class index 1) probability per sample, used for AUC.
        let mut positive_scores = Array1::zeros(num_samples);
        // Accumulated squared error of the positive-class probability.
        let mut loss_sum = 0.0;

        // Add extra metrics fields that will be populated later
        let mut class_accuracies = vec![0.0; self.class_labels.len()];
        let class_labels = self.class_labels.clone();

        for i in 0..num_samples {
            let features = x_test.row(i).to_owned();
            let probabilities = self.predict_proba(&features)?;

            // Predicted label is the arg-max of the class probabilities.
            let mut pred_idx = 0usize;
            for k in 1..probabilities.len() {
                if probabilities[k] > probabilities[pred_idx] {
                    pred_idx = k;
                }
            }

            y_pred[i] = pred_idx as f64;
            confidences[i] = probabilities[pred_idx];

            // Positive-class score (binary problems use class index 1).
            let positive_score = if probabilities.len() >= 2 {
                probabilities[1]
            } else {
                probabilities[0]
            };
            positive_scores[i] = positive_score;

            let target = if y_test[i] > 0.5 { 1.0 } else { 0.0 };
            let diff = positive_score - target;
            loss_sum += diff * diff;
        }

        // Compute metrics
        let mut tp = 0.0;
        let mut fp = 0.0;
        let mut tn = 0.0;
        let mut fn_ = 0.0;

        for i in 0..num_samples {
            let true_label = y_test[i];
            let pred_label = y_pred[i];

            // Binary classification metrics
            if true_label > 0.5 {
                if pred_label > 0.5 {
                    tp += 1.0;
                } else {
                    fn_ += 1.0;
                }
            } else {
                if pred_label > 0.5 {
                    fp += 1.0;
                } else {
                    tn += 1.0;
                }
            }
        }

        let accuracy = (tp + tn) / num_samples as f64;

        let precision = if tp + fp > 0.0 { tp / (tp + fp) } else { 0.0 };

        let recall = if tp + fn_ > 0.0 { tp / (tp + fn_) } else { 0.0 };

        let f1_score = if precision + recall > 0.0 {
            2.0 * precision * recall / (precision + recall)
        } else {
            0.0
        };

        // Real rank-based (Mann-Whitney) AUC from the positive-class scores.
        let auc = compute_binary_auc(&positive_scores, y_test);
        let confusion_matrix =
            Array2::from_shape_vec((2, 2), vec![tn, fp, fn_, tp]).map_err(|e| {
                MLError::MLOperationError(format!("Failed to create confusion matrix: {}", e))
            })?;

        // Calculate per-class accuracies
        for (i, label) in self.class_labels.iter().enumerate() {
            let class_samples = y_test
                .iter()
                .enumerate()
                .filter(|(_, &y)| y == i as f64)
                .map(|(idx, _)| idx)
                .collect::<Vec<_>>();

            if !class_samples.is_empty() {
                let correct = class_samples
                    .iter()
                    .filter(|&&idx| y_pred[idx] == i as f64)
                    .count();

                class_accuracies[i] = correct as f64 / class_samples.len() as f64;
            }
        }

        // Return metrics with the added fields
        Ok(ClassificationMetrics {
            accuracy,
            precision,
            recall,
            f1_score,
            auc,
            confusion_matrix,
            class_accuracies,
            class_labels,
            average_loss: loss_sum / num_samples as f64,
        })
    }

    /// Returns the class-probability distribution for a sample.
    ///
    /// The (fixed or trained) QNN is evaluated on `features`; its per-class
    /// Pauli-Z expectation values are converted to a probability distribution
    /// with a numerically stable soft-max.
    pub fn predict_proba(&self, features: &Array1<f64>) -> Result<Array1<f64>> {
        let logits = self.qnn.forward(features)?;
        if logits.is_empty() {
            return Err(MLError::MLOperationError(
                "QNN produced an empty output for HEP prediction".to_string(),
            ));
        }

        let max_logit = logits.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        let mut probabilities: Vec<f64> = logits.iter().map(|&v| (v - max_logit).exp()).collect();
        let sum: f64 = probabilities.iter().sum();
        if sum > 0.0 {
            for value in probabilities.iter_mut() {
                *value /= sum;
            }
        }
        Ok(Array1::from_vec(probabilities))
    }

    /// Predicts the class for a sample using the quantum neural network.
    ///
    /// Returns the highest-probability class label and its probability
    /// (confidence).
    pub fn predict(&self, features: &Array1<f64>) -> Result<(String, f64)> {
        let probabilities = self.predict_proba(features)?;

        let mut best_idx = 0usize;
        for k in 1..probabilities.len() {
            if probabilities[k] > probabilities[best_idx] {
                best_idx = k;
            }
        }

        if best_idx < self.class_labels.len() {
            Ok((self.class_labels[best_idx].clone(), probabilities[best_idx]))
        } else {
            Err(MLError::MLOperationError(format!(
                "Invalid prediction index: {}",
                best_idx
            )))
        }
    }

    /// Computes feature importance
    pub fn feature_importance(&self) -> Result<Array1<f64>> {
        // In a real implementation, this would compute feature importance
        // through perturbation analysis or gradient-based methods
        let mut importance = Array1::zeros(self.feature_dimension);

        for i in 0..self.feature_dimension {
            importance[i] = thread_rng().random::<f64>();
        }

        // Normalize
        let sum = importance.sum();
        if sum > 0.0 {
            importance.mapv_inplace(|x| x / sum);
        }

        Ok(importance)
    }
}

/// Rank-based (Mann-Whitney U) area under the ROC curve for a binary problem.
///
/// `scores` are the model's positive-class scores and `labels` are the true
/// binary labels (positive when `> 0.5`).  Ties in the scores receive their
/// average rank.  When either class is absent the metric is undefined and the
/// chance value `0.5` is returned.
fn compute_binary_auc(scores: &Array1<f64>, labels: &Array1<f64>) -> f64 {
    let n = scores.len();
    if n == 0 {
        return 0.5;
    }

    // Sort sample indices by ascending score.
    let mut order: Vec<usize> = (0..n).collect();
    order.sort_by(|&a, &b| {
        scores[a]
            .partial_cmp(&scores[b])
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    // Assign 1-based ranks, averaging ranks across tied score groups.
    let mut ranks = vec![0.0_f64; n];
    let mut i = 0;
    while i < n {
        let mut j = i;
        while j + 1 < n && (scores[order[j + 1]] - scores[order[i]]).abs() < 1e-12 {
            j += 1;
        }
        let average_rank = ((i + 1) + (j + 1)) as f64 / 2.0;
        for &idx in &order[i..=j] {
            ranks[idx] = average_rank;
        }
        i = j + 1;
    }

    let mut sum_positive_ranks = 0.0;
    let mut n_positive = 0.0;
    let mut n_negative = 0.0;
    for i in 0..n {
        if labels[i] > 0.5 {
            sum_positive_ranks += ranks[i];
            n_positive += 1.0;
        } else {
            n_negative += 1.0;
        }
    }

    if n_positive == 0.0 || n_negative == 0.0 {
        return 0.5;
    }

    (sum_positive_ranks - n_positive * (n_positive + 1.0) / 2.0) / (n_positive * n_negative)
}

/// Specialized detector for Higgs bosons in collision data
#[derive(Debug, Clone)]
pub struct HiggsDetector {
    /// Quantum neural network
    qnn: QuantumNeuralNetwork,

    /// Number of qubits
    num_qubits: usize,
}

impl HiggsDetector {
    /// Creates a new Higgs detector
    pub fn new(num_qubits: usize) -> Result<Self> {
        // Create a QNN for Higgs detection
        let layers = vec![
            crate::qnn::QNNLayerType::EncodingLayer { num_features: 10 },
            crate::qnn::QNNLayerType::VariationalLayer {
                num_params: 2 * num_qubits,
            },
            crate::qnn::QNNLayerType::EntanglementLayer {
                connectivity: "full".to_string(),
            },
            crate::qnn::QNNLayerType::VariationalLayer {
                num_params: 2 * num_qubits,
            },
            crate::qnn::QNNLayerType::MeasurementLayer {
                measurement_basis: "computational".to_string(),
            },
        ];

        let qnn = QuantumNeuralNetwork::new(
            layers, num_qubits, 10, // Input dimension
            1,  // Output dimension (binary)
        )?;

        Ok(HiggsDetector { qnn, num_qubits })
    }

    /// Detects Higgs bosons in a collision event
    pub fn detect_higgs(&self, event: &CollisionEvent) -> Result<Vec<bool>> {
        // For each particle, predict whether it's a Higgs boson
        let mut results = Vec::with_capacity(event.particles.len());

        for particle in &event.particles {
            let score = self.score_particle(particle)?;
            results.push(score > 0.7); // Threshold for Higgs detection
        }

        Ok(results)
    }

    /// Computes a score for a particle (higher = more likely to be a Higgs)
    pub fn score_particle(&self, particle: &ParticleFeatures) -> Result<f64> {
        // Dummy implementation
        match particle.particle_type {
            ParticleType::Higgs => Ok(0.85 + 0.15 * thread_rng().random::<f64>()),
            _ => Ok(0.2 * thread_rng().random::<f64>()),
        }
    }
}

/// System for detecting particle collision anomalies
#[derive(Debug, Clone)]
pub struct ParticleCollisionClassifier {
    qnn: QuantumNeuralNetwork,
    num_qubits: usize,
}

impl ParticleCollisionClassifier {
    /// Creates a new particle collision classifier
    pub fn new() -> Self {
        // This is a placeholder implementation
        let layers = vec![
            crate::qnn::QNNLayerType::EncodingLayer { num_features: 10 },
            crate::qnn::QNNLayerType::VariationalLayer { num_params: 20 },
            crate::qnn::QNNLayerType::EntanglementLayer {
                connectivity: "full".to_string(),
            },
            crate::qnn::QNNLayerType::MeasurementLayer {
                measurement_basis: "computational".to_string(),
            },
        ];

        let qnn = QuantumNeuralNetwork::new(
            layers, 8,  // 8 qubits
            10, // 10 features
            2,  // 2 classes
        )
        .expect("should create ParticleCollisionClassifier QNN");

        ParticleCollisionClassifier { qnn, num_qubits: 8 }
    }

    /// Builder method to set the number of qubits
    pub fn with_qubits(mut self, num_qubits: usize) -> Self {
        self.num_qubits = num_qubits;
        self
    }

    /// Builder method to set the feature dimension
    pub fn with_input_features(self, _features: usize) -> Self {
        // This would normally update the QNN, but we'll just return self for now
        self
    }

    /// Builder method to set the number of measurement qubits
    pub fn with_measurement_qubits(self, _num_qubits: usize) -> Result<Self> {
        // This would normally update the QNN, but we'll just return self for now
        Ok(self)
    }

    /// Trains the classifier
    pub fn train(
        &mut self,
        data: &Array2<f64>,
        labels: &Array1<f64>,
        epochs: usize,
        learning_rate: f64,
    ) -> Result<crate::qnn::TrainingResult> {
        self.qnn.train_1d(data, labels, epochs, learning_rate)
    }

    /// Evaluates the classifier
    pub fn evaluate(
        &self,
        data: &Array2<f64>,
        labels: &Array1<f64>,
    ) -> Result<ClassificationMetrics> {
        // Dummy implementation
        Ok(ClassificationMetrics {
            accuracy: 0.85,
            precision: 0.82,
            recall: 0.88,
            f1_score: 0.85,
            auc: 0.91,
            confusion_matrix: Array2::eye(2),
            class_accuracies: vec![0.85, 0.86], // Dummy class accuracies
            class_labels: vec!["Signal".to_string(), "Background".to_string()], // Dummy class labels
            average_loss: 0.15, // Dummy average loss
        })
    }
}

/// Event reconstructor for HEP data
#[derive(Debug, Clone)]
pub struct EventReconstructor {
    qnn: QuantumNeuralNetwork,
    input_dim: usize,
    output_dim: usize,
}

impl EventReconstructor {
    /// Creates a new event reconstructor
    pub fn new() -> Self {
        // This is a placeholder implementation
        let layers = vec![
            crate::qnn::QNNLayerType::EncodingLayer { num_features: 10 },
            crate::qnn::QNNLayerType::VariationalLayer { num_params: 20 },
            crate::qnn::QNNLayerType::EntanglementLayer {
                connectivity: "full".to_string(),
            },
            crate::qnn::QNNLayerType::MeasurementLayer {
                measurement_basis: "computational".to_string(),
            },
        ];

        let qnn = QuantumNeuralNetwork::new(
            layers, 8,  // 8 qubits
            10, // 10 input features
            10, // 10 output features
        )
        .expect("should create EventReconstructor QNN");

        EventReconstructor {
            qnn,
            input_dim: 10,
            output_dim: 10,
        }
    }

    /// Builder method to set the input dimension
    pub fn with_input_features(mut self, input_dim: usize) -> Self {
        self.input_dim = input_dim;
        self
    }

    /// Builder method to set the output dimension
    pub fn with_output_features(mut self, output_dim: usize) -> Self {
        self.output_dim = output_dim;
        self
    }

    /// Builder method to set the number of quantum layers
    pub fn with_quantum_layers(self, _num_layers: usize) -> Result<Self> {
        // This would normally update the QNN, but we'll just return self for now
        Ok(self)
    }
}

/// Anomaly detector for HEP data
#[derive(Debug, Clone)]
pub struct AnomalyDetector {
    features: usize,
    quantum_encoder: bool,
}

impl AnomalyDetector {
    /// Creates a new anomaly detector
    pub fn new() -> Self {
        AnomalyDetector {
            features: 10,
            quantum_encoder: false,
        }
    }

    /// Builder method to set the number of features
    pub fn with_features(mut self, features: usize) -> Self {
        self.features = features;
        self
    }

    /// Builder method to enable/disable quantum encoding
    pub fn with_quantum_encoder(mut self, quantum_encoder: bool) -> Self {
        self.quantum_encoder = quantum_encoder;
        self
    }

    /// Builder method to set the kernel method
    pub fn with_kernel_method(self, _kernel_method: KernelMethod) -> Result<Self> {
        // This would normally update the anomaly detector, but we'll just return self for now
        Ok(self)
    }
}

/// Kernel method for quantum machine learning
#[derive(Debug, Clone, Copy)]
pub enum KernelMethod {
    /// Linear kernel
    Linear,

    /// Polynomial kernel
    Polynomial,

    /// Quantum kernel
    QuantumKernel,
}