Support PAULI-SUM Gate Specifications

Cargo.toml

@@ -7,3 +7,6 @@ members = [
 [profile.release]
 lto = true
 codegen-units = 1
+
+[workspace.dependencies]
+strum = { version = "0.24.1", features = ["derive"] }

quil-py/Cargo.toml

@@ -17,6 +17,7 @@ crate-type = ["cdylib", "rlib"]
 
 [dependencies]
 quil-rs = { path = "../quil-rs", version="0.16.0-rc.0" }
+strum.workspace = true
 # pyo3 dependencies should be updated together
 pyo3 = { version = "0.17" }
 rigetti-pyo3 = "0.1.0-rc.6"

quil-py/quil/instructions/__init__.pyi

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Dict, List, Optional, Union, final
+from typing import Dict, List, Optional, Tuple, Union, final
 
 from quil.expression import Expression
 
@@ -437,7 +437,6 @@ class GateModifier(Enum):
     Forked = "FORKED"
 
 class Gate:
-    @classmethod
     def __new__(
         cls,
         name: str,
@@ -471,6 +470,48 @@ class Gate:
         """
         ...
 
+@final
+class PauliGate(Enum):
+    I = "I"
+    X = "X"
+    Y = "Y"
+    Z = "Z"
+    @staticmethod
+    def parse(word: str) -> "PauliGate":
+        """
+        Parses a ``PauliGate`` from a string. Raises a ``ParseEnumError`` if the
+        string isn't a valid Pauli word.
+        """
+        ...
+
+class PauliTerm:
+    @staticmethod
+    def __new__(
+        cls,
+        arguments: List[Tuple[PauliGate, str]],
+        expression: Expression,
+    ) -> "PauliTerm": ...
+    @property
+    def arguments(self) -> List[Tuple[PauliGate, str]]: ...
+    @arguments.setter
+    def arguments(self, word: List[Tuple[PauliGate, str]]): ...
+    @property
+    def expression(self) -> Expression: ...
+    @expression.setter
+    def expression(self, expression: Expression): ...
+
+class PauliSum:
+    @staticmethod
+    def __new__(cls, arguments: List[str], terms: List[PauliTerm]) -> "PauliSum": ...
+    @property
+    def arguments(self) -> List[str]: ...
+    @arguments.setter
+    def arguments(self, arguments: List[str]): ...
+    @property
+    def terms(self) -> List[PauliTerm]: ...
+    @terms.setter
+    def terms(self, terms: List[PauliTerm]): ...
+
 @final
 class GateSpecification:
     """
@@ -487,24 +528,28 @@ class GateSpecification:
         - from_*: Creates a new ``GateSpecification`` using an instance of the inner type for the variant.
     """
 
-    def inner(self) -> Union[List[List[Expression]], List[int]]:
+    def inner(self) -> Union[List[List[Expression]], List[int], PauliSum]:
         """
         Returns the inner value of the variant. Raises a ``RuntimeError`` if inner data doesn't exist.
         """
         ...
     def is_matrix(self) -> bool: ...
     def is_permutation(self) -> bool: ...
+    def is_pauli_sum(self) -> bool: ...
     def as_matrix(self) -> Optional[List[List[Expression]]]: ...
     def to_matrix(self) -> List[List[Expression]]: ...
     def as_permutation(self) -> Optional[List[int]]: ...
     def to_permutation(self) -> List[int]: ...
+    def as_pauli_sum(self) -> Optional[PauliSum]: ...
+    def to_pauli_sum(self) -> PauliSum: ...
     @staticmethod
     def from_matrix(matrix: List[List[Expression]]) -> "GateSpecification": ...
     @staticmethod
     def from_permutation(permutation: List[int]) -> "GateSpecification": ...
+    @staticmethod
+    def from_pauli_sum(pauli_term: PauliSum) -> "GateSpecification": ...
 
 class GateDefinition:
-    @classmethod
     def __new__(
         cls,
         name: str,
@@ -513,10 +558,16 @@ class GateDefinition:
     ) -> "GateDefinition": ...
     @property
     def name(self) -> str: ...
+    @name.setter
+    def name(self, name: str): ...
     @property
     def parameters(self) -> List[str]: ...
+    @parameters.setter
+    def parameters(self, parameters: List[str]): ...
     @property
     def specification(self) -> GateSpecification: ...
+    @specification.setter
+    def specification(self, specification: GateSpecification): ...
 
 @final
 class Qubit:

quil-py/src/expression.rs

@@ -33,6 +33,7 @@ py_wrap_error!(
 
 py_wrap_union_enum! {
     #[derive(Debug, Hash, PartialEq, Eq)]
+    #[pyo3(module="quil.expression")]
     PyExpression(Expression) as "Expression" {
         address: Address => PyMemoryReference,
         function_call: FunctionCall => PyFunctionCallExpression,

quil-py/src/instruction/gate.rs

@@ -1,26 +1,32 @@
 use quil_rs::{
     expression::Expression,
-    instruction::{Gate, GateDefinition, GateModifier, GateSpecification, Qubit},
+    instruction::{
+        Gate, GateDefinition, GateModifier, GateSpecification, PauliGate, PauliSum, PauliTerm,
+        Qubit,
+    },
 };
+use strum;
 
 use rigetti_pyo3::{
-    impl_repr, impl_str, py_wrap_data_struct, py_wrap_error, py_wrap_simple_enum,
-    py_wrap_union_enum,
+    impl_from_str, impl_parse, impl_repr, impl_str, py_wrap_data_struct, py_wrap_error,
+    py_wrap_simple_enum, py_wrap_type, py_wrap_union_enum,
     pyo3::{
         exceptions::PyValueError,
         pyclass::CompareOp,
         pymethods,
         types::{PyInt, PyString},
-        IntoPy, Py, PyObject, PyResult, Python,
+        IntoPy, Py, PyErr, PyObject, PyResult, Python,
     },
-    wrap_error, PyTryFrom, PyWrapper, ToPython, ToPythonError,
+    wrap_error, PyTryFrom, PyWrapper, PyWrapperMut, ToPython, ToPythonError,
 };
 
 use super::PyQubit;
 use crate::expression::PyExpression;
 
 wrap_error!(RustGateError(quil_rs::instruction::GateError));
 py_wrap_error!(quil, RustGateError, GateError, PyValueError);
+wrap_error!(RustParseEnumError(strum::ParseError));
+py_wrap_error!(quil, RustParseEnumError, EnumParseError, PyValueError);
 
 py_wrap_data_struct! {
     #[derive(Debug, PartialEq, Eq)]
@@ -114,11 +120,137 @@ impl PyGateModifier {
     }
 }
 
+py_wrap_simple_enum! {
+    #[derive(Debug, PartialEq, Eq)]
+    PyPauliGate(PauliGate) as "PauliGate" {
+        I,
+        X,
+        Y,
+        Z
+    }
+}
+impl_repr!(PyPauliGate);
+impl_str!(PyPauliGate);
+impl_from_str!(PyPauliGate, RustParseEnumError);
+impl_parse!(PyPauliGate);
+
+// This is a helper type to help manage easy conversion of the inner tuple
+// with the macros. It should not be exposed directly.
+py_wrap_type! {
+    PyPauliPair((PauliGate, String))
+}
+
+impl PyPauliPair {
+    pub(crate) fn from_py_tuple(py: Python<'_>, tuple: (PyPauliGate, String)) -> PyResult<Self> {
+        Ok(Self((PauliGate::py_try_from(py, &tuple.0)?, tuple.1)))
+    }
+}
+
+py_wrap_data_struct! {
+    #[derive(Debug, PartialEq, Eq)]
+    #[pyo3(subclass)]
+    PyPauliTerm(PauliTerm) as "PauliTerm" {
+        arguments: Vec<(PauliGate, String)> => Vec<PyPauliPair>,
+        expression: Expression => PyExpression
+    }
+}
+
+#[pymethods]
+impl PyPauliTerm {
+    #[new]
+    pub fn new(
+        py: Python<'_>,
+        arguments: Vec<(PyPauliGate, String)>,
+        expression: PyExpression,
+    ) -> PyResult<Self> {
+        Ok(Self(PauliTerm::new(
+            Vec::<(PauliGate, String)>::py_try_from(
+                py,
+                &PyPauliTerm::py_pairs_from_tuples(py, arguments)?,
+            )?,
+            Expression::py_try_from(py, &expression)?,
+        )))
+    }
+
+    // Override the getters/setters generated by [`py_wrap_data_struct!`] so that they
+    // return/take tuples instead of the wrapping [`PyPauliPair`] type.
+    #[getter(arguments)]
+    fn get_arguments_as_tuple(&self, py: Python<'_>) -> PyResult<Vec<(PyPauliGate, String)>> {
+        let mut pairs: Vec<(PyPauliGate, String)> =
+            Vec::with_capacity(self.as_inner().arguments.len());
+        self.as_inner()
+            .arguments
+            .iter()
+            .try_for_each(|(gate, arg)| {
+                pairs.push((gate.to_python(py)?, arg.clone()));
+                Ok::<(), PyErr>(())
+            })?;
+        Ok(pairs)
+    }
+
+    #[setter(arguments)]
+    fn set_arguments_from_tuple(
+        &mut self,
+        py: Python<'_>,
+        arguments: Vec<(PyPauliGate, String)>,
+    ) -> PyResult<()> {
+        self.as_inner_mut().arguments = Vec::<(PauliGate, String)>::py_try_from(
+            py,
+            &PyPauliTerm::py_pairs_from_tuples(py, arguments)?,
+        )?;
+        Ok(())
+    }
+}
+
+impl PyPauliTerm {
+    pub(crate) fn py_pairs_from_tuples(
+        py: Python<'_>,
+        tuples: Vec<(PyPauliGate, String)>,
+    ) -> PyResult<Vec<PyPauliPair>> {
+        let mut pairs: Vec<PyPauliPair> = Vec::with_capacity(tuples.len());
+        tuples.into_iter().try_for_each(|tuple| {
+            pairs.push(PyPauliPair::from_py_tuple(py, tuple)?);
+            Ok::<(), PyErr>(())
+        })?;
+        Ok(pairs)
+    }
+}
+
+py_wrap_data_struct! {
+    #[derive(Debug, PartialEq, Eq)]
+    #[pyo3(subclass)]
+    PyPauliSum(PauliSum) as "PauliSum" {
+        arguments: Vec<String> => Vec<Py<PyString>>,
+        terms: Vec<PauliTerm> => Vec<PyPauliTerm>
+    }
+}
+impl_repr!(PyPauliSum);
+
+#[pymethods]
+impl PyPauliSum {
+    #[new]
+    pub fn new(py: Python<'_>, arguments: Vec<String>, terms: Vec<PyPauliTerm>) -> PyResult<Self> {
+        Ok(Self(
+            PauliSum::new(arguments, Vec::<PauliTerm>::py_try_from(py, &terms)?)
+                .map_err(RustGateError::from)
+                .map_err(RustGateError::to_py_err)?,
+        ))
+    }
+
+    pub fn __richcmp__(&self, py: Python<'_>, other: &Self, op: CompareOp) -> PyObject {
+        match op {
+            CompareOp::Eq => (self.as_inner() == other.as_inner()).into_py(py),
+            _ => py.NotImplemented(),
+        }
+    }
+}
+
 py_wrap_union_enum! {
     #[derive(Debug, PartialEq, Eq)]
     PyGateSpecification(GateSpecification) as "GateSpecification" {
         matrix: Matrix => Vec<Vec<PyExpression>>,
-        permutation: Permutation => Vec<Py<PyInt>>
+        permutation: Permutation => Vec<Py<PyInt>>,
+        pauli_sum: PauliSum => PyPauliSum
     }
 }
 impl_repr!(PyGateSpecification);
@@ -155,11 +287,15 @@ impl PyGateDefinition {
         parameters: Vec<String>,
         specification: PyGateSpecification,
     ) -> PyResult<Self> {
-        Ok(Self(GateDefinition::new(
-            name,
-            parameters,
-            GateSpecification::py_try_from(py, &specification)?,
-        )))
+        Ok(Self(
+            GateDefinition::new(
+                name,
+                parameters,
+                GateSpecification::py_try_from(py, &specification)?,
+            )
+            .map_err(RustGateError::from)
+            .map_err(RustGateError::to_py_err)?,
+        ))
     }
 
     pub fn __richcmp__(&self, py: Python<'_>, other: &Self, op: CompareOp) -> PyObject {

quil-py/src/instruction/mod.rs

@@ -16,7 +16,10 @@ pub use self::{
         PySharing, PyVector,
     },
     frame::{PyAttributeValue, PyFrameAttributes, PyFrameDefinition, PyFrameIdentifier},
-    gate::{GateError, PyGate, PyGateDefinition, PyGateModifier, PyGateSpecification},
+    gate::{
+        GateError, PyGate, PyGateDefinition, PyGateModifier, PyGateSpecification, PyPauliGate,
+        PyPauliSum, PyPauliTerm,
+    },
     measurement::PyMeasurement,
     qubit::PyQubit,
     waveform::{PyWaveform, PyWaveformDefinition, PyWaveformInvocation},
@@ -89,6 +92,9 @@ create_init_submodule! {
         PyGateDefinition,
         PyGateModifier,
         PyGateSpecification,
+        PyPauliGate,
+        PyPauliTerm,
+        PyPauliSum,
         PyMeasurement,
         PyMemoryReference,
         PyQubit,

quil-py/test/instructions/test_gate.py

@@ -0,0 +1,10 @@
+from quil.expression import Expression
+from quil.instructions import PauliTerm, PauliGate
+
+
+class TestPauliTerm:
+    def test_new(self):
+        pt = PauliTerm([(PauliGate.X, "a")], Expression.new_pi())
+        assert pt.arguments == [(PauliGate.X, "a")]
+        pt.arguments = [(PauliGate.Y, "b")]
+        assert pt.arguments == [(PauliGate.Y, "b")]

quil-rs/Cargo.toml

@@ -18,7 +18,7 @@ nom_locate = "4.0.0"
 num-complex = "0.4.0"
 petgraph = "0.6.2"
 serde = { version = "1.0.125", features = ["derive"] }
-strum = { version = "0.24.1", features = ["derive"] }
+strum.workspace = true
 thiserror = "1.0.37"
 once_cell = "1.17.1"