burn_quantum_visualization 0.3.0

Interactive Bevy workbench for quantum circuits, states, photonic fabrics and optimization
Documentation
use super::*;
use crate::simulation::{read_tensor, snapshot_pure};
use burn::tensor::{DType, Tensor, TensorData};
use burn_quantum::{
    DenseSimulator, StateVector, optim::qn_spsa_sample, photonic::runtime::shot_seed,
};

pub(super) async fn build(device: &Device, limits: TraceLimits) -> Result<VisualizationDocument> {
    const UPDATES: usize = 24;
    let mut circuit = Circuit::new(Register::qubits(1)?, 1, 0);
    circuit.rotate(0, Rotation::qubit(Axis::Y, Angle::parameter(0)))?;
    limits.check(circuit.register(), UPDATES + 1)?;
    let engine = DenseSimulator::new(limits.resources());
    let prepare = |angle: f64| -> Result<StateVector> {
        circuit.run(
            &engine,
            Some(Tensor::from_data(
                TensorData::new(vec![angle], [1]),
                (device, DType::F32),
            )),
            None,
            1,
            device,
        )
    };
    let mut theta = 0.35;
    let mut metric = 0.25;
    let mut state = prepare(theta)?;
    let mut frame = snapshot_pure(&state, 0, "Initial parameter".into(), limits).await?;
    let point =
        |iteration: usize, theta: f64, metric: f64, frame: &TimelineFrame| OptimizationPoint {
            iteration,
            parameters: vec![theta],
            loss: frame.probabilities[0] - frame.probabilities[1],
            metric: vec![metric],
            bloch: frame.subsystems.iter().map(|s| s.bloch).collect(),
            probabilities: frame.probabilities.clone(),
            loss_evaluations: 1 + 3 * iteration,
            fidelity_evaluations: 4 * iteration,
        };
    let mut points = vec![point(0, theta, metric, &frame)];
    let mut frames = vec![frame];
    for iteration in 0..UPDATES {
        let eps = 0.08;
        let d1 = if shot_seed(731, iteration as u64, 0) & 1 == 0 {
            1.
        } else {
            -1.
        };
        let d2 = if shot_seed(731, iteration as u64, 1) & 1 == 0 {
            1.
        } else {
            -1.
        };
        let plus = prepare(theta + eps * d1)?;
        let minus = prepare(theta - eps * d1)?;
        let plus_prob = read_tensor(plus.probabilities()).await?;
        let minus_prob = read_tensor(minus.probabilities()).await?;
        let both = prepare(theta + eps * (d1 + d2))?;
        let crossed = prepare(theta + eps * (-d1 + d2))?;
        let fidelities = [
            read_tensor(state.fidelity(&plus)?).await?[0],
            read_tensor(state.fidelity(&minus)?).await?[0],
            read_tensor(state.fidelity(&both)?).await?[0],
            read_tensor(state.fidelity(&crossed)?).await?[0],
        ];
        let sample = qn_spsa_sample(
            eps,
            &[d1],
            &[d2],
            [plus_prob[0] - plus_prob[1], minus_prob[0] - minus_prob[1]],
            fidelities,
            1,
        )?;
        metric += (sample.metric[0] - metric) / (iteration + 2) as f64;
        theta -= 0.12 * sample.gradient[0] / (metric.abs() + 0.05);
        state = prepare(theta)?;
        frame = snapshot_pure(
            &state,
            iteration + 1,
            format!("QN-SPSA update {}", iteration + 1),
            limits,
        )
        .await?;
        points.push(point(iteration + 1, theta, metric, &frame));
        frames.push(frame);
    }
    let kind = DemoKind::Optimization;
    let document = VisualizationDocument {
        schema_version: 1,
        title: kind.title().into(),
        description: kind.description().into(),
        circuit,
        parameters: vec![theta],
        inputs: vec![],
        operation_labels: vec!["Ry(theta)".into()],
        trace: ExperimentTrace {
            kind: TraceKind::Optimization,
            frames,
            backend: format!("{device:?}"),
        },
        fabric: None,
        optimization: points,
        protocol: None,
    };
    document.validate_with(limits)?;
    Ok(document)
}