Skip to main content

prism_q/circuit/
parameter.rs

1//! Parameter slots over a circuit's rotation angles, shared by gradient and
2//! binding consumers.
3
4use std::mem::Discriminant;
5
6use super::{Circuit, Instruction, SmallVec};
7use crate::error::{PrismError, Result};
8use crate::gates::Gate;
9
10/// Binds one rotation gate instruction to a slot in the parameter vector.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub struct ParamLink {
13    /// Index into [`Circuit::instructions`].
14    pub instruction: usize,
15    /// Slot in the parameter vector this gate reads.
16    pub slot: usize,
17}
18
19/// Parameter slots over the rotation angles of a circuit.
20///
21/// A slot may drive several instructions (weight sharing): [`bind`](Self::bind)
22/// writes one angle to each, and the adjoint gradient accumulates their
23/// contributions into one entry. Arity is declared at construction rather than
24/// inferred from the links, so a value vector of the wrong length is rejected
25/// even when the trailing slots carry no link.
26///
27/// Bindable gates are `Rx`, `Ry`, `Rz`, `Rzz`, `P`, and `PauliRot`, the `Gate`
28/// variants carrying a rotation angle.
29#[derive(Debug, Clone, Default, PartialEq, Eq)]
30pub struct Parameters {
31    links: Vec<ParamLink>,
32    num_slots: usize,
33    /// Slot names, empty when the set is positional only. OpenQASM `input`
34    /// declarations are named, so export and the parser share these.
35    names: Vec<String>,
36    /// Gate kind and targets at each linked index when the set was built.
37    /// Links are instruction indices, so an edited circuit would otherwise
38    /// rebind silently; [`validate`](Parameters::validate) compares this.
39    shape: Vec<(Discriminant<Gate>, SmallVec<[usize; 4]>)>,
40}
41
42impl Parameters {
43    /// Declare `num_slots` slots with no links yet.
44    pub fn new(num_slots: usize) -> Self {
45        Self {
46            num_slots,
47            ..Default::default()
48        }
49    }
50
51    /// Build from explicit links over a vector of `num_slots` entries.
52    pub fn from_links(links: Vec<ParamLink>, num_slots: usize) -> Self {
53        Self {
54            links,
55            num_slots,
56            ..Default::default()
57        }
58    }
59
60    /// Name the slots, in slot order. Names are what OpenQASM `input`
61    /// declarations carry and what export emits.
62    ///
63    /// # Panics
64    /// Panics if `names` is not exactly [`num_slots`](Self::num_slots) long.
65    pub fn with_names<I, S>(mut self, names: I) -> Self
66    where
67        I: IntoIterator<Item = S>,
68        S: Into<String>,
69    {
70        let names: Vec<String> = names.into_iter().map(Into::into).collect();
71        assert_eq!(
72            names.len(),
73            self.num_slots,
74            "expected {} slot names, got {}",
75            self.num_slots,
76            names.len()
77        );
78        self.names = names;
79        self
80    }
81
82    /// Name of `slot`, or `None` when the set is positional only.
83    pub fn name_of(&self, slot: usize) -> Option<&str> {
84        self.names.get(slot).map(String::as_str)
85    }
86
87    /// Slot a name refers to.
88    pub fn slot_of(&self, name: &str) -> Option<usize> {
89        self.names.iter().position(|n| n == name)
90    }
91
92    /// Record the gate kind and targets each link points at, so a later
93    /// [`validate`](Self::validate) against an edited circuit fails loudly
94    /// rather than binding the wrong gates.
95    pub fn pinned_to(mut self, circuit: &Circuit) -> Self {
96        self.shape = self
97            .links
98            .iter()
99            .filter_map(|link| match circuit.instructions.get(link.instruction) {
100                Some(Instruction::Gate { gate, targets }) => {
101                    Some((std::mem::discriminant(gate), targets.clone()))
102                }
103                _ => None,
104            })
105            .collect();
106        if self.shape.len() != self.links.len() {
107            self.shape.clear();
108        }
109        self
110    }
111
112    /// Give every bindable gate its own slot, in circuit order. The common
113    /// case for a variational ansatz where each rotation is independent.
114    pub fn all_rotations(circuit: &Circuit) -> Self {
115        let mut links = Vec::new();
116        for (i, inst) in circuit.instructions.iter().enumerate() {
117            if let Instruction::Gate { gate, .. } = inst {
118                if gate.pauli_generator().is_some() {
119                    links.push(ParamLink {
120                        instruction: i,
121                        slot: links.len(),
122                    });
123                }
124            }
125        }
126        let num_slots = links.len();
127        Self {
128            links,
129            num_slots,
130            ..Default::default()
131        }
132        .pinned_to(circuit)
133    }
134
135    /// Record that `instruction` reads `slot`, widening the declared slot count
136    /// to cover it. For builders accumulating links before the count is known.
137    pub(super) fn link_growing(&mut self, instruction: usize, slot: usize) {
138        self.num_slots = self.num_slots.max(slot + 1);
139        self.links.push(ParamLink { instruction, slot });
140    }
141
142    /// Record that `instruction` reads `slot`.
143    ///
144    /// # Panics
145    /// Panics if `slot` is not below the declared slot count. Slot bounds are
146    /// fixed when the set is constructed, so an out-of-range slot is a caller
147    /// bug rather than bad input.
148    pub fn link(&mut self, instruction: usize, slot: usize) {
149        assert!(
150            slot < self.num_slots,
151            "slot {} out of bounds (parameter set declares {} slots)",
152            slot,
153            self.num_slots
154        );
155        self.links.push(ParamLink { instruction, slot });
156    }
157
158    pub fn links(&self) -> &[ParamLink] {
159        &self.links
160    }
161
162    /// Length of the value vector [`bind`](Self::bind) expects.
163    pub fn num_slots(&self) -> usize {
164        self.num_slots
165    }
166
167    /// True when no instruction is linked. A set may still declare slots.
168    pub fn is_empty(&self) -> bool {
169        self.links.is_empty()
170    }
171
172    /// Check every link against `circuit` without binding.
173    ///
174    /// A declared slot that no instruction reads is accepted: OpenQASM allows an
175    /// `input` the body never uses, and a sweep may range over a superset of the
176    /// circuit's parameters. Use [`unread_slots`](Self::unread_slots) to report
177    /// them.
178    ///
179    /// # Errors
180    /// Returns [`PrismError::InvalidParameter`] when a link points past the end
181    /// of the instruction stream, at a non-gate instruction, or at a gate
182    /// carrying no angle.
183    pub fn validate(&self, circuit: &Circuit) -> Result<()> {
184        let n = circuit.instructions.len();
185        for link in &self.links {
186            if link.instruction >= n {
187                return Err(PrismError::InvalidParameter {
188                    message: format!(
189                        "parameter link references instruction {} but the circuit has {n} instructions",
190                        link.instruction
191                    ),
192                });
193            }
194            match &circuit.instructions[link.instruction] {
195                Instruction::Gate { gate, .. } if gate.pauli_generator().is_some() => {}
196                Instruction::Gate { gate, .. } => {
197                    return Err(PrismError::InvalidParameter {
198                        message: format!(
199                            "instruction {} (`{}`) carries no bindable angle; bindable gates are rx, ry, rz, rzz, p, pauli_rot",
200                            link.instruction,
201                            gate.name()
202                        ),
203                    });
204                }
205                _ => {
206                    return Err(PrismError::InvalidParameter {
207                        message: format!(
208                            "parameter link references instruction {} which is not a gate",
209                            link.instruction
210                        ),
211                    });
212                }
213            }
214        }
215
216        if !self.shape.is_empty() {
217            for (link, (kind, targets)) in self.links.iter().zip(&self.shape) {
218                let Instruction::Gate { gate, targets: at } =
219                    &circuit.instructions[link.instruction]
220                else {
221                    unreachable!("link validated as a gate above")
222                };
223                if std::mem::discriminant(gate) != *kind || at.as_slice() != targets.as_slice() {
224                    return Err(PrismError::InvalidParameter {
225                        message: format!(
226                            "instruction {} no longer holds the gate this parameter set was built against; the circuit was edited after the links were recorded",
227                            link.instruction
228                        ),
229                    });
230                }
231            }
232        }
233
234        Ok(())
235    }
236
237    /// Declared slots that no instruction reads, whose bound values are
238    /// discarded. Legal, but usually a mistake worth surfacing.
239    pub fn unread_slots(&self) -> Vec<usize> {
240        let mut used = vec![false; self.num_slots];
241        for link in &self.links {
242            used[link.slot] = true;
243        }
244        (0..self.num_slots).filter(|s| !used[*s]).collect()
245    }
246
247    /// Write `values` into a copy of `template`.
248    ///
249    /// # Errors
250    /// Returns [`PrismError::InvalidParameter`] when `values.len()` differs from
251    /// [`num_slots`](Self::num_slots), when any value is not finite, or when
252    /// [`validate`](Self::validate) rejects the links.
253    pub fn bind(&self, template: &Circuit, values: &[f64]) -> Result<Circuit> {
254        let mut out = template.clone();
255        self.bind_into(template, values, &mut out)?;
256        Ok(out)
257    }
258
259    /// Bind into an existing circuit, reusing its allocations.
260    ///
261    /// `out` is overwritten with `template` and then patched, so a sweep can
262    /// hold one buffer across every point instead of allocating per binding.
263    ///
264    /// # Errors
265    /// Same conditions as [`bind`](Self::bind).
266    pub fn bind_into(&self, template: &Circuit, values: &[f64], out: &mut Circuit) -> Result<()> {
267        self.check_values(values)?;
268        self.validate(template)?;
269
270        out.num_qubits = template.num_qubits;
271        out.num_classical_bits = template.num_classical_bits;
272        out.instructions.clone_from(&template.instructions);
273        self.write_angles(out, values);
274        Ok(())
275    }
276
277    /// Check `values` against the declared arity without touching a circuit.
278    ///
279    /// # Errors
280    /// Same arity and finiteness conditions as [`bind`](Self::bind); the link
281    /// checks belong to [`validate`](Self::validate) and are not repeated here.
282    pub(crate) fn check_values(&self, values: &[f64]) -> Result<()> {
283        if values.len() != self.num_slots {
284            return Err(PrismError::InvalidParameter {
285                message: format!(
286                    "expected {} parameter values, got {}",
287                    self.num_slots,
288                    values.len()
289                ),
290            });
291        }
292        if let Some(i) = values.iter().position(|v| !v.is_finite()) {
293            return Err(PrismError::InvalidParameter {
294                message: format!(
295                    "parameter value {i} is {}, expected a finite angle",
296                    values[i]
297                ),
298            });
299        }
300        Ok(())
301    }
302
303    /// Overwrite the linked angles of a circuit already shaped like the
304    /// template. Callers must have validated the links against it.
305    pub(crate) fn write_angles(&self, out: &mut Circuit, values: &[f64]) {
306        for link in &self.links {
307            *angle_mut(&mut out.instructions[link.instruction]) = values[link.slot];
308        }
309    }
310
311    /// Read the angle each slot currently holds in `circuit`.
312    ///
313    /// # Errors
314    /// Same conditions as [`validate`](Self::validate), which this runs first
315    /// because a link pointing at a non-angle gate has no value to read.
316    pub fn values(&self, circuit: &Circuit) -> Result<Vec<f64>> {
317        self.validate(circuit)?;
318        let mut out = vec![0.0; self.num_slots];
319        for link in &self.links {
320            out[link.slot] = angle_of(&circuit.instructions[link.instruction]);
321        }
322        Ok(out)
323    }
324}
325
326/// Angle of a gate already validated as bindable.
327pub(crate) fn angle_of(instruction: &Instruction) -> f64 {
328    match instruction {
329        Instruction::Gate {
330            gate: Gate::Rx(t) | Gate::Ry(t) | Gate::Rz(t) | Gate::Rzz(t) | Gate::P(t),
331            ..
332        } => *t,
333        Instruction::Gate {
334            gate: Gate::PauliRot(data),
335            ..
336        } => data.theta(),
337        _ => unreachable!("parameter link validated as bindable"),
338    }
339}
340
341pub(crate) fn angle_mut(instruction: &mut Instruction) -> &mut f64 {
342    match instruction {
343        Instruction::Gate {
344            gate: Gate::Rx(t) | Gate::Ry(t) | Gate::Rz(t) | Gate::Rzz(t) | Gate::P(t),
345            ..
346        } => t,
347        Instruction::Gate {
348            gate: Gate::PauliRot(data),
349            ..
350        } => &mut data.theta,
351        _ => unreachable!("parameter link validated as bindable"),
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    fn two_rotations() -> Circuit {
360        let mut c = Circuit::new(2, 0);
361        c.add_gate(Gate::Rx(0.1), &[0]);
362        c.add_gate(Gate::Cx, &[0, 1]);
363        c.add_gate(Gate::Rz(0.2), &[1]);
364        c
365    }
366
367    #[test]
368    fn all_rotations_declares_one_slot_per_gate() {
369        let p = Parameters::all_rotations(&two_rotations());
370        assert_eq!(p.num_slots(), 2);
371        assert_eq!(p.links().len(), 2);
372        assert_eq!(p.links()[1].instruction, 2);
373    }
374
375    #[test]
376    fn bind_writes_angles_and_leaves_structure() {
377        let template = two_rotations();
378        let p = Parameters::all_rotations(&template);
379        let bound = p.bind(&template, &[1.5, 2.5]).unwrap();
380        assert_eq!(bound.instructions.len(), 3);
381        assert!(matches!(
382            bound.instructions[0],
383            Instruction::Gate {
384                gate: Gate::Rx(t),
385                ..
386            } if t == 1.5
387        ));
388        assert!(matches!(
389            bound.instructions[2],
390            Instruction::Gate {
391                gate: Gate::Rz(t),
392                ..
393            } if t == 2.5
394        ));
395    }
396
397    #[test]
398    fn shared_slot_writes_every_linked_gate() {
399        let template = two_rotations();
400        let mut p = Parameters::new(1);
401        p.link(0, 0);
402        p.link(2, 0);
403        let bound = p.bind(&template, &[0.75]).unwrap();
404        assert_eq!(super::angle_of(&bound.instructions[0]), 0.75);
405        assert_eq!(super::angle_of(&bound.instructions[2]), 0.75);
406    }
407
408    #[test]
409    fn wrong_arity_is_an_error() {
410        let template = two_rotations();
411        let p = Parameters::all_rotations(&template);
412        assert!(p.bind(&template, &[1.0]).is_err());
413        assert!(p.bind(&template, &[1.0, 2.0, 3.0]).is_err());
414    }
415
416    #[test]
417    fn non_finite_value_is_an_error() {
418        let template = two_rotations();
419        let p = Parameters::all_rotations(&template);
420        assert!(p.bind(&template, &[f64::NAN, 0.0]).is_err());
421        assert!(p.bind(&template, &[0.0, f64::INFINITY]).is_err());
422    }
423
424    #[test]
425    fn link_past_end_is_an_error() {
426        let template = two_rotations();
427        let mut p = Parameters::new(1);
428        p.link(99, 0);
429        assert!(p.bind(&template, &[0.5]).is_err());
430    }
431
432    #[test]
433    fn link_to_non_bindable_gate_is_an_error() {
434        let template = two_rotations();
435        let mut p = Parameters::new(1);
436        p.link(1, 0);
437        assert!(p.bind(&template, &[0.5]).is_err());
438    }
439
440    #[test]
441    fn slot_no_gate_reads_is_accepted_and_reported() {
442        let template = two_rotations();
443        let mut p = Parameters::new(2);
444        p.link(0, 0);
445        assert!(p.bind(&template, &[0.5, 0.5]).is_ok());
446        assert_eq!(p.unread_slots(), vec![1]);
447    }
448
449    #[test]
450    #[should_panic(expected = "out of bounds")]
451    fn slot_past_declared_count_panics() {
452        let mut p = Parameters::new(1);
453        p.link(0, 4);
454    }
455
456    #[test]
457    fn values_round_trip_through_bind() {
458        let template = two_rotations();
459        let p = Parameters::all_rotations(&template);
460        let bound = p.bind(&template, &[0.3, 0.4]).unwrap();
461        assert_eq!(p.values(&bound).unwrap(), vec![0.3, 0.4]);
462    }
463
464    // `pauli_generator` decides which gates bind, and `angle_of`, `angle_mut`,
465    // and the replay recipe's `write_angle` each match the same variants by
466    // hand. A new angle-carrying variant fails here until all four agree.
467    // `BatchRzz` is the one deliberate exception: replay writes its edges
468    // although the fused gate itself is not bindable.
469    #[test]
470    fn angle_sites_agree_on_every_gate() {
471        use crate::circuit::plan::write_angle;
472        use crate::gates::PauliRotData;
473        use crate::sim::unified_pauli::PauliAxis;
474
475        let h = Gate::H.matrix_2x2();
476        let gates = [
477            Gate::Rx(0.1),
478            Gate::Ry(0.2),
479            Gate::Rz(0.3),
480            Gate::P(0.4),
481            Gate::Rzz(0.5),
482            Gate::PauliRot(Box::new(PauliRotData {
483                theta: 0.6,
484                axes: vec![PauliAxis::X, PauliAxis::Z],
485            })),
486            Gate::Id,
487            Gate::H,
488            Gate::T,
489            Gate::SX,
490            Gate::Cx,
491            Gate::Cz,
492            Gate::Swap,
493            Gate::Cu(Box::new(h)),
494            Gate::Fused(Box::new(h)),
495            Gate::Fused2q(Box::new(Gate::Cx.matrix_4x4())),
496            Gate::QftBlock { start: 0, num: 2 },
497        ];
498        for gate in gates {
499            let targets: SmallVec<[usize; 4]> = (0..gate.num_qubits()).collect();
500            let mut inst = Instruction::Gate {
501                gate: gate.clone(),
502                targets,
503            };
504            let bindable = gate.pauli_generator().is_some();
505            assert_eq!(write_angle(&mut inst, 0, 1.25), bindable, "{gate}");
506            if bindable {
507                assert_eq!(angle_of(&inst), 1.25, "{gate}");
508                *angle_mut(&mut inst) = 2.5;
509                assert_eq!(angle_of(&inst), 2.5, "{gate}");
510            }
511        }
512    }
513}