rustiq_core/routines/
rotation_optimization.rs

1use crate::structures::pauli_dag::{build_dag_from_pauli_set, get_front_layer, Dag};
2/// Implementation of the Zhang et al algorithm for T-count optimization
3/// The algorithm is adapted to work with any angles (not just pi/4 or pi/2)
4use crate::structures::{CliffordCircuit, CliffordGate, Parameter, PauliLike, PauliSet, Tableau};
5use std::collections::HashSet;
6
7fn update_rot_pi2<T: PauliLike>(axis: &str, k: i32, rest: &mut T, dagger: bool) {
8    let support: Vec<_> = (0..axis.len())
9        .filter(|i| axis.chars().nth(*i).unwrap() != 'I')
10        .collect();
11    for qbit in support.iter() {
12        match axis.chars().nth(*qbit).unwrap() {
13            'X' => rest.h(*qbit),
14            'Y' => {
15                if dagger {
16                    rest.sqrt_xd(*qbit)
17                } else {
18                    rest.sqrt_x(*qbit)
19                }
20            }
21            _ => {}
22        }
23    }
24    let target = support[0];
25    for qbit in support.iter().skip(1) {
26        rest.cnot(*qbit, target);
27    }
28    for _ in 0..k {
29        if dagger {
30            rest.sd(target)
31        } else {
32            rest.s(target)
33        }
34    }
35    for qbit in support.iter().skip(1) {
36        rest.cnot(*qbit, target);
37    }
38    for qbit in support.iter() {
39        match axis.chars().nth(*qbit).unwrap() {
40            'X' => rest.h(*qbit),
41            'Y' => {
42                if dagger {
43                    rest.sqrt_x(*qbit)
44                } else {
45                    rest.sqrt_xd(*qbit)
46                }
47            }
48            _ => {}
49        }
50    }
51}
52
53fn zhang_internal(
54    rotations: &[(String, Parameter)],
55    nqubits: usize,
56    inverse_final_clifford: &mut Tableau,
57) -> Vec<(String, Parameter)> {
58    let mut bucket = PauliSet::new(nqubits);
59    let mut axes = Vec::new();
60    let mut future_angles = Vec::new();
61    for (axis, angle) in rotations.iter() {
62        axes.push(axis.clone());
63        future_angles.push(angle.clone());
64    }
65    let mut rest = PauliSet::from_slice(&axes);
66    let mut angles: Vec<Parameter> = Vec::new();
67    let mut rot_index = 0;
68    while !rest.is_empty() {
69        let (phase, axis) = rest.get(0);
70        rest.pop();
71        let angle = &mut future_angles[rot_index];
72        if phase {
73            angle.flip_sign();
74        }
75        let index = bucket.insert(&axis, false);
76        let mut merged = false;
77        for i in (0..index).rev() {
78            if !bucket.commute(index, i) {
79                break;
80            }
81            if bucket.equals(i, index) {
82                angles[i] += angle.clone();
83                merged = true;
84                let (new_angle, mult_pi_2) = angles[i].simplify();
85                update_rot_pi2(&axis, mult_pi_2, &mut rest, true);
86                update_rot_pi2(&axis, mult_pi_2, inverse_final_clifford, true);
87                angles[i] = new_angle;
88                if angles[i].is_zero_mod_two_pi() {
89                    bucket.set_to_identity(i);
90                }
91                break;
92            }
93        }
94        if merged {
95            bucket.pop_last();
96        } else {
97            angles.push(angle.clone());
98        }
99        rot_index += 1;
100    }
101    let mut output = Vec::new();
102    for (i, angle) in angles.iter().enumerate().take(bucket.len()) {
103        let (phase, pstring) = bucket.get(i);
104        assert!(!phase);
105        if pstring.chars().any(|c| c != 'I') {
106            output.push((pstring, angle.clone()));
107        }
108    }
109    output
110}
111
112/// Implementation of the Zhang et al algorithm for T-count optimization adapted to work
113/// with any angles (even parametrized rotations)
114pub fn zhang_rotation_optimization(
115    mut rotations: Vec<(String, Parameter)>,
116    nqubits: usize,
117) -> (Vec<(String, Parameter)>, Tableau) {
118    let mut current_size = rotations.len();
119    let mut inverse_final_clifford = Tableau::new(nqubits);
120    loop {
121        let new_rotations = zhang_internal(&rotations, nqubits, &mut inverse_final_clifford);
122        if new_rotations.len() == current_size {
123            break;
124        }
125        current_size = new_rotations.len();
126        rotations = new_rotations;
127    }
128    (rotations, inverse_final_clifford)
129}
130
131/// Data structure of the initial state propagation
132/// It stores a Pauli DAG and a set of marked qubits (i.e. qubits that are not stabilized by |0> anymore)
133struct MarkedPauliDag {
134    pauli_set: PauliSet,
135    final_clifford: Tableau,
136    dag: Dag,
137    marked: HashSet<usize>,
138    output_pauli_set: PauliSet,
139    output_rotations: Vec<usize>,
140    did_something: bool,
141}
142
143impl MarkedPauliDag {
144    fn new(pauli_set: PauliSet) -> Self {
145        let dag = build_dag_from_pauli_set(&pauli_set);
146        Self {
147            output_pauli_set: PauliSet::new(pauli_set.n),
148            final_clifford: Tableau::new(pauli_set.n),
149            pauli_set,
150            dag,
151            marked: HashSet::new(),
152            output_rotations: Vec::new(),
153            did_something: false,
154        }
155    }
156
157    fn get_front_indices(&self) -> Vec<usize> {
158        let front_layer = get_front_layer(&self.dag);
159        front_layer
160            .into_iter()
161            .map(|ni| *self.dag.node_weight(ni).unwrap())
162            .collect()
163    }
164
165    fn get_unmarked_xy(&self, rotation_index: usize) -> Option<usize> {
166        (0..self.pauli_set.n).find(|&qbit| {
167            self.pauli_set.get_entry(qbit, rotation_index) & !self.marked.contains(&qbit)
168        })
169    }
170    fn has_unmarked_xy(&self, rotation_index: usize) -> bool {
171        for qbit in 0..self.pauli_set.n {
172            if self.pauli_set.get_entry(qbit, rotation_index) & !self.marked.contains(&qbit) {
173                return true;
174            }
175        }
176        false
177    }
178    fn get_rotation_score(&self, rotation_index: usize) -> usize {
179        if self.has_unmarked_xy(rotation_index) {
180            return self.pauli_set.support_size(rotation_index) - 1;
181        }
182        let mut score = 0;
183        for qbit in 0..self.pauli_set.n {
184            if !self.pauli_set.get_entry(qbit, rotation_index)
185                & self
186                    .pauli_set
187                    .get_entry(qbit + self.pauli_set.n, rotation_index)
188                & !self.marked.contains(&qbit)
189            {
190                score += 1;
191            }
192        }
193        score
194    }
195    /// Optimize a given rotation from the front layer of the PDAG
196    fn optimize_rotation(&mut self, rotation_index: usize) {
197        // Remove unmarked Z components:
198        for qbit in 0..self.pauli_set.n {
199            if !self.marked.contains(&qbit)
200                && self
201                    .pauli_set
202                    .get_entry(qbit + self.pauli_set.n, rotation_index)
203                    & !self.pauli_set.get_entry(qbit, rotation_index)
204            {
205                self.pauli_set.set_entry(rotation_index, qbit, false, false);
206                self.did_something = true;
207            }
208        }
209        // Fold X components onto a single unmarked qubit (if there is any)
210        let mut piece = CliffordCircuit::new(self.pauli_set.n);
211        if let Some(control) = self.get_unmarked_xy(rotation_index) {
212            for qbit in self.pauli_set.get_support(rotation_index) {
213                if qbit == control {
214                    continue;
215                }
216                assert!(
217                    self.marked.contains(&qbit) | self.pauli_set.get_entry(qbit, rotation_index)
218                );
219                if self
220                    .pauli_set
221                    .get_entry(qbit + self.pauli_set.n, rotation_index)
222                {
223                    if self.pauli_set.get_entry(qbit, rotation_index) {
224                        piece.gates.push(CliffordGate::Sd(qbit));
225                    } else {
226                        piece.gates.push(CliffordGate::H(qbit));
227                    }
228                }
229                piece.gates.push(CliffordGate::CNOT(control, qbit));
230            }
231            for qbit in self.pauli_set.get_support(rotation_index) {
232                if qbit == control {
233                    continue;
234                }
235                if self
236                    .pauli_set
237                    .get_entry(qbit + self.pauli_set.n, rotation_index)
238                {
239                    if self.pauli_set.get_entry(qbit, rotation_index) {
240                        piece.gates.push(CliffordGate::S(qbit));
241                    } else {
242                        piece.gates.push(CliffordGate::H(qbit));
243                    }
244                }
245                self.pauli_set.set_entry(rotation_index, qbit, false, false);
246                self.did_something = true;
247            }
248            self.marked.insert(control);
249        }
250        if self.pauli_set.support_size(rotation_index) > 0 {
251            let (phase, bv) = self.pauli_set.get_as_vec_bool(rotation_index);
252            self.output_pauli_set.insert_vec_bool(&bv, phase);
253            self.pauli_set.conjugate_with_circuit(&piece.dagger());
254            let c = Tableau::from_circuit(&piece);
255            self.final_clifford = self.final_clifford.clone() * c;
256            self.output_rotations.push(rotation_index);
257        }
258        self.dag.retain_nodes(|graph, node_index| {
259            *graph.node_weight(node_index).unwrap() != rotation_index
260        });
261    }
262
263    /// Attempts to simplify a rotation from the front layer of the PDAG
264    /// Returns true if a simplification was made
265    fn simplify_once(&mut self) -> bool {
266        let front_layer = self.get_front_indices();
267        // Looking for the best rotation to remove
268        // score(R) = |R| -1 if there is an unmarked X/Y
269        // score(R) = # of unmarked Zs otherwise
270        let best_candidate = front_layer
271            .into_iter()
272            .map(|ri| (ri, self.get_rotation_score(ri)))
273            .max_by_key(|(_, score)| *score);
274        if let Some((rotation_index, score)) = best_candidate {
275            if score > 0 {
276                self.optimize_rotation(rotation_index);
277                return true;
278            }
279        }
280        false
281    }
282
283    fn pop_rest(&mut self) {
284        loop {
285            let front_layer = self.get_front_indices();
286            if front_layer.is_empty() {
287                break;
288            }
289            for rotation_index in front_layer.iter() {
290                let (phase, bv) = self.pauli_set.get_as_vec_bool(*rotation_index);
291                self.output_pauli_set.insert_vec_bool(&bv, phase);
292                self.output_rotations.push(*rotation_index);
293            }
294            self.dag.retain_nodes(|graph, node_index| {
295                !front_layer.contains(graph.node_weight(node_index).unwrap())
296            })
297        }
298    }
299
300    pub fn propagate(mut self) -> (PauliSet, Tableau, Vec<usize>, bool) {
301        while self.simplify_once() {}
302        self.pop_rest();
303        (
304            self.output_pauli_set,
305            self.final_clifford,
306            self.output_rotations,
307            self.did_something,
308        )
309    }
310}
311
312pub fn full_initial_state_propagation(
313    rotations: &[(String, Parameter)],
314) -> (Vec<(String, Parameter)>, Tableau) {
315    let axes: Vec<_> = rotations.iter().map(|e| e.0.clone()).collect();
316    let mut angles: Vec<_> = rotations.iter().map(|e| e.1.clone()).collect();
317    let mut pset = PauliSet::from_slice(&axes);
318    let mut final_clifford = Tableau::new(pset.n);
319    loop {
320        let mpdag = MarkedPauliDag::new(pset);
321        let (new_pset, clifford, new_rotations, carry_on) = mpdag.propagate();
322        angles = new_rotations
323            .into_iter()
324            .map(|i| angles[i].clone())
325            .collect();
326        pset = new_pset;
327        final_clifford = final_clifford * clifford;
328        if !carry_on {
329            break;
330        }
331    }
332    let mut new_rotations = Vec::new();
333    for (i, angle) in angles.iter_mut().enumerate().take(pset.len()) {
334        let (phase, string) = pset.get(i);
335        if phase {
336            angle.flip_sign();
337        }
338        new_rotations.push((string, angle.clone()));
339    }
340    (new_rotations, final_clifford)
341}
rustiq_core/routines/rotation_optimization.rs

rustiq_core/routines/
rotation_optimization.rs
