quil_rs/instruction/
waveform.rs

1use indexmap::IndexMap;
2
3#[cfg(not(feature = "python"))]
4use optipy::strip_pyo3;
5#[cfg(feature = "stubs")]
6use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
7
8use crate::{
9    expression::Expression,
10    pickleable_new,
11    quil::{write_join_quil, Quil, INDENT},
12};
13
14use super::write_parameter_string;
15
16#[derive(Clone, Debug, PartialEq, Eq, Hash)]
17#[cfg_attr(feature = "stubs", gen_stub_pyclass)]
18#[cfg_attr(
19    feature = "python",
20    pyo3::pyclass(module = "quil.instructions", eq, frozen, hash, get_all, subclass)
21)]
22pub struct Waveform {
23    pub matrix: Vec<Expression>,
24    pub parameters: Vec<String>,
25}
26
27#[cfg_attr(feature = "stubs", gen_stub_pymethods)]
28#[cfg_attr(feature = "python", pyo3::pymethods)]
29#[cfg_attr(not(feature = "python"), strip_pyo3)]
30impl Waveform {
31    #[new]
32    pub fn new(matrix: Vec<Expression>, parameters: Vec<String>) -> Self {
33        Self { matrix, parameters }
34    }
35}
36
37#[derive(Clone, Debug, PartialEq, Eq, Hash)]
38#[cfg_attr(feature = "stubs", gen_stub_pyclass)]
39#[cfg_attr(
40    feature = "python",
41    pyo3::pyclass(module = "quil.instructions", eq, frozen, hash, get_all, subclass)
42)]
43pub struct WaveformDefinition {
44    pub name: String,
45    pub definition: Waveform,
46}
47
48pickleable_new! {
49    impl WaveformDefinition {
50        pub fn new(name: String, definition: Waveform);
51    }
52}
53
54impl Quil for WaveformDefinition {
55    fn write(
56        &self,
57        f: &mut impl std::fmt::Write,
58        fall_back_to_debug: bool,
59    ) -> crate::quil::ToQuilResult<()> {
60        write!(f, "DEFWAVEFORM {}", self.name)?;
61        write_parameter_string(f, &self.definition.parameters)?;
62        write!(f, ":\n{INDENT}")?;
63        write_join_quil(f, fall_back_to_debug, &self.definition.matrix, ", ", "")
64    }
65}
66
67#[cfg(test)]
68mod test_waveform_definition {
69    use super::{Waveform, WaveformDefinition};
70    use crate::expression::Expression;
71    use crate::quil::Quil;
72    use crate::real;
73
74    use insta::assert_snapshot;
75    use rstest::rstest;
76
77    #[rstest]
78    #[case(
79        "Simple WaveformDefinition",
80        WaveformDefinition {
81            name: "WF".to_string(),
82            definition: Waveform {
83                matrix: vec![
84                    Expression::Number(real!(0.0)),
85                    Expression::Number(real!(0.0)),
86                    Expression::Number(real!(0.00027685415721916584))
87                ],
88                parameters: vec!["theta".to_string()]
89            }
90        }
91    )]
92    fn test_display(#[case] description: &str, #[case] waveform_def: WaveformDefinition) {
93        insta::with_settings!({
94            snapshot_suffix => description,
95        }, {
96            assert_snapshot!(waveform_def.to_quil_or_debug())
97        })
98    }
99}
100
101pub type WaveformParameters = IndexMap<String, Expression>;
102
103#[derive(Clone, Debug, PartialEq, Eq)]
104#[cfg_attr(feature = "stubs", gen_stub_pyclass)]
105#[cfg_attr(
106    feature = "python",
107    pyo3::pyclass(module = "quil.instructions", eq, frozen, get_all, subclass)
108)]
109pub struct WaveformInvocation {
110    pub name: String,
111    pub parameters: WaveformParameters,
112}
113
114pickleable_new! {
115    impl WaveformInvocation {
116        pub fn new(name: String, parameters: WaveformParameters);
117    }
118}
119
120impl Quil for WaveformInvocation {
121    fn write(
122        &self,
123        f: &mut impl std::fmt::Write,
124        _fall_back_to_debug: bool,
125    ) -> crate::quil::ToQuilResult<()> {
126        let mut key_value_pairs = self
127            .parameters
128            .iter()
129            .collect::<Vec<(&String, &Expression)>>();
130
131        key_value_pairs.sort_by(|(k1, _), (k2, _)| k1.cmp(k2));
132
133        if key_value_pairs.is_empty() {
134            write!(f, "{}", self.name,)?;
135        } else {
136            write!(
137                f,
138                "{}({})",
139                self.name,
140                key_value_pairs
141                    .iter()
142                    .map(|(k, v)| format!("{k}: {}", v.to_quil_or_debug()))
143                    .collect::<Vec<String>>()
144                    .join(", ")
145            )?;
146        }
147        Ok(())
148    }
149}
150
151#[cfg(test)]
152mod waveform_invocation_tests {
153    use super::WaveformParameters;
154
155    use crate::{instruction::WaveformInvocation, quil::Quil};
156
157    #[test]
158    fn format_no_parameters() {
159        let invocation = WaveformInvocation {
160            name: "CZ".into(),
161            parameters: WaveformParameters::new(),
162        };
163        assert_eq!(invocation.to_quil_or_debug(), "CZ".to_string());
164    }
165}