paulistrings 0.1.2

//! [`PauliString<W>`], the symplectic `(x, z)` encoding of a Pauli operator on up to `64·W` qubits.
//! The derived `Ord`, lexicographic in `x` then `z`, must agree with the `(x, z)` tuple order every bucket of a [`PauliSum`](crate::PauliSum) is sorted in.

use num_complex::Complex64;

use crate::phase::Phase;

/// A Pauli operator on up to `64·W` qubits, `W` being the number of 64-bit words per `x`/`z` part.
///
/// Pick the smallest `W` that fits; the Python bindings monomorphize `W ∈ {1, 2, 4, 8, 16}`.
/// The phase of a product is not stored: [`Self::mul_assign`] returns it as a [`Phase`] for the caller to fold into a coefficient.
/// `Y` is the key `(x=1, z=1)` with no phase factor.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
#[repr(C)]
pub struct PauliString<const W: usize> {
    /// X-part bitmask: bit `q` set iff qubit `q` is `X` or `Y`.
    pub x: [u64; W],
    /// Z-part bitmask: bit `q` set iff qubit `q` is `Z` or `Y`.
    pub z: [u64; W],
}

impl<const W: usize> PauliString<W> {
    /// Identity Pauli string (all qubits `I`).
    pub const fn identity() -> Self {
        Self { x: [0u64; W], z: [0u64; W] }
    }

    /// Single-qubit `X` Pauli on `qubit`.
    ///
    /// # Panics
    ///
    /// Panics in debug builds if `qubit >= 64 · W`.
    pub fn x(qubit: usize) -> Self {
        debug_assert!(qubit < 64 * W);
        let mut p = Self::identity();
        p.x[qubit / 64] = 1u64 << (qubit % 64);
        p
    }

    /// Single-qubit `Y = (x=1, z=1)` on `qubit`; the `i` in `Y = i·X·Z` is the caller's concern.
    ///
    /// # Panics
    ///
    /// Panics in debug builds if `qubit >= 64 · W`.
    pub fn y(qubit: usize) -> Self {
        debug_assert!(qubit < 64 * W);
        let mut p = Self::identity();
        let word = qubit / 64;
        let bit = 1u64 << (qubit % 64);
        p.x[word] = bit;
        p.z[word] = bit;
        p
    }

    /// Single-qubit `Z` Pauli on `qubit`.
    ///
    /// # Panics
    ///
    /// Panics in debug builds if `qubit >= 64 · W`.
    pub fn z(qubit: usize) -> Self {
        debug_assert!(qubit < 64 * W);
        let mut p = Self::identity();
        p.z[qubit / 64] = 1u64 << (qubit % 64);
        p
    }

    /// Number of non-identity qubits.
    pub fn weight(&self) -> u32 {
        (0..W).map(|i| (self.x[i] | self.z[i]).count_ones()).sum()
    }

    /// Multiply `self * other` in place, returning the phase `i^k` such that the true product is `phase · self`.
    pub fn mul_assign(&mut self, other: &Self) -> Phase {
        // Per qubit, P(a,b)·P(c,d) = i^δ·P(a⊕c, b⊕d) with δ = 2bc + ab + cd − (a⊕c)(b⊕d) mod 4, from P(a,b) = i^{ab}·X^a·Z^b.
        let mut delta: u32 = 0;
        for i in 0..W {
            let a = self.x[i];
            let b = self.z[i];
            let c = other.x[i];
            let d = other.z[i];
            let zc_x = (b & c).count_ones();
            let y_self = (a & b).count_ones();
            let y_other = (c & d).count_ones();
            let y_result = ((a ^ c) & (b ^ d)).count_ones();
            delta = delta.wrapping_add(zc_x.wrapping_mul(2));
            delta = delta.wrapping_add(y_self);
            delta = delta.wrapping_add(y_other);
            delta = delta.wrapping_sub(y_result);
            self.x[i] ^= c;
            self.z[i] ^= d;
        }
        Phase::new(delta as u8)
    }

    /// Value-returning multiply: `(self * other, phase)`.
    pub fn product(mut self, other: &Self) -> (Self, Phase) {
        let phase = self.mul_assign(other);
        (self, phase)
    }

    /// `true` iff every set bit lies on a qubit index `< num_qubits`.
    ///
    /// # Panics
    ///
    /// Panics in debug builds if `num_qubits > 64 · W`.
    pub fn is_within(&self, num_qubits: usize) -> bool {
        debug_assert!(num_qubits <= 64 * W);
        let mut leak: u64 = 0;
        for i in 0..W {
            leak |= (self.x[i] | self.z[i]) & !word_mask(num_qubits, i);
        }
        leak == 0
    }

    /// `true` iff `self` and `other` commute as Pauli operators.
    pub fn commutes_with(&self, other: &Self) -> bool {
        // Parity is GF(2)-linear, so XOR-fold the words and take one popcount.
        let mut folded: u64 = 0;
        for i in 0..W {
            folded ^= (self.x[i] & other.z[i]) ^ (self.z[i] & other.x[i]);
        }
        folded.count_ones() & 1 == 0
    }

    /// Commutator `[self, other] = self·other − other·self`, as `(product, coefficient)`.
    ///
    /// The coefficient is `2·i^k` when the strings anticommute and exactly `0` when they commute; the returned string is [`Self::product`]'s result either way.
    pub fn commutator(self, other: &Self) -> (Self, Complex64) {
        let vanishes = self.commutes_with(other);
        let (product, phase) = self.product(other);
        (product, scaled_phase(phase, vanishes))
    }

    /// Anticommutator `{self, other} = self·other + other·self`, as `(product, coefficient)`.
    ///
    /// The coefficient is `2·i^k` when the strings commute and exactly `0` when they anticommute.
    pub fn anticommutator(self, other: &Self) -> (Self, Complex64) {
        let vanishes = !self.commutes_with(other);
        let (product, phase) = self.product(other);
        (product, scaled_phase(phase, vanishes))
    }
}

/// Mask of the live qubit bits in word `word`, given `num_qubits` total.
pub(crate) fn word_mask(num_qubits: usize, word: usize) -> u64 {
    let first_qubit = 64 * word;
    if num_qubits >= first_qubit + 64 {
        !0u64
    } else if num_qubits <= first_qubit {
        0
    } else {
        (1u64 << (num_qubits - first_qubit)) - 1
    }
}

/// The bit mask of `qubits`, each asserted below `num_qubits`.
pub(crate) fn qubit_mask<const W: usize>(
    qubits: impl IntoIterator<Item = usize>,
    num_qubits: usize,
) -> [u64; W] {
    let mut mask = [0u64; W];
    for q in qubits {
        assert!(q < num_qubits, "qubit {q} is outside the {num_qubits}-qubit register");
        mask[q / 64] |= 1u64 << (q % 64);
    }
    mask
}

/// `0` when the bracket vanishes, `2·i^k` otherwise.
fn scaled_phase(phase: Phase, vanishes: bool) -> Complex64 {
    if vanishes {
        Complex64::new(0.0, 0.0)
    } else {
        phase.to_complex() * 2.0
    }
}

impl<const W: usize> Default for PauliString<W> {
    fn default() -> Self {
        Self::identity()
    }
}

#[cfg(test)]
mod tests;