Add type annotations to the PauliSum class

CHANGELOG.md

@@ -16,6 +16,7 @@ Changelog
 -   Added support for the `XY` (parameterized `iSWAP`) gate family in `Program`s
     and in `ISA`s (@ecpeterson, gh-1096, gh-1107, gh-1111).
 -   Removed the `tox.ini` and `readthedocs.yml` files (@karalekas, gh-1108).
+-   Type hints have been added to the `PauliSum` class (@rht, gh-1104).
 
 ### Bugfixes
 

pyquil/paulis.py

@@ -22,17 +22,19 @@
 import numpy as np
 import copy
 
-from typing import Dict, FrozenSet, Iterable, Iterator, List, Optional, Tuple, Union
+from typing import Callable, Dict, FrozenSet, Iterable, Iterator, List, Optional, Sequence, Tuple, Union
 
 from pyquil.quilatom import QubitPlaceholder
 
 from .quil import Program
 from .gates import H, RZ, RX, CNOT, X, PHASE, QUANTUM_GATES
 from numbers import Number
 from collections import OrderedDict
-from collections.abc import Sequence
 import warnings
 
+PauliCoefficientDesignator = Union[int, float, complex]
+PauliDesignator = Union['PauliTerm', 'PauliSum']
+
 PAULI_OPS = ["X", "Y", "Z", "I"]
 PAULI_PROD = {'ZZ': 'I', 'YY': 'I', 'XX': 'I', 'II': 'I',
               'XY': 'Z', 'XZ': 'Y', 'YX': 'Z', 'YZ': 'X', 'ZX': 'Y',
@@ -61,7 +63,7 @@ def __init__(self, *args, **kwargs):
 """
 
 
-def _valid_qubit(index):
+def _valid_qubit(index: int) -> bool:
     return ((isinstance(index, integer_types) and index >= 0)
             or isinstance(index, QubitPlaceholder))
 
@@ -70,7 +72,7 @@ class PauliTerm(object):
     """A term is a product of Pauli operators operating on different qubits.
     """
 
-    def __init__(self, op: str, index: int, coefficient: Union[int, float, complex] = 1.0):
+    def __init__(self, op: str, index: int, coefficient: PauliCoefficientDesignator = 1.0):
         """ Create a new Pauli Term with a Pauli operator at a particular index and a leading
         coefficient.
 
@@ -203,7 +205,7 @@ def _multiply_factor(self, factor: str, index: int) -> 'PauliTerm':
 
         return new_term
 
-    def __mul__(self, term: Union['PauliTerm', 'PauliSum', Number]) -> Union['PauliTerm', 'PauliSum']:
+    def __mul__(self, term: Union[PauliDesignator, PauliCoefficientDesignator]) -> PauliDesignator:
         """Multiplies this Pauli Term with another PauliTerm, PauliSum, or number according to the
         Pauli algebra rules.
 
@@ -253,7 +255,7 @@ def __pow__(self, power: int) -> 'PauliTerm':
             result *= self
         return result
 
-    def __add__(self, other: Union['PauliTerm', 'PauliSum', Number]) -> 'PauliSum':
+    def __add__(self, other: Union[PauliDesignator, PauliCoefficientDesignator]) -> 'PauliSum':
         """Adds this PauliTerm with another one.
 
         :param other: A PauliTerm object, a PauliSum object, or a Number
@@ -502,7 +504,7 @@ class PauliSum(object):
     """A sum of one or more PauliTerms.
     """
 
-    def __init__(self, terms):
+    def __init__(self, terms: Sequence[PauliTerm]):
         """
         :param Sequence terms: A Sequence of PauliTerms.
         """
@@ -514,7 +516,7 @@ def __init__(self, terms):
         else:
             self.terms = terms
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         """Equality testing to see if two PauliSum's are equivalent.
 
         :param PauliSum other: The PauliSum to compare this PauliSum with.
@@ -531,30 +533,30 @@ def __eq__(self, other):
 
         return set(self.terms) == set(other.terms)
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash(frozenset(self.terms))
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return " + ".join([str(term) for term in self.terms])
 
-    def __len__(self):
+    def __len__(self) -> int:
         """
         The length of the PauliSum is the number of PauliTerms in the sum.
         """
         return len(self.terms)
 
-    def __getitem__(self, item):
+    def __getitem__(self, item: int) -> PauliTerm:
         """
         :param int item: The index of the term in the sum to return
         :return: The PauliTerm at the index-th position in the PauliSum
         :rtype: PauliTerm
         """
         return self.terms[item]
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[PauliTerm]:
         return self.terms.__iter__()
 
-    def __mul__(self, other):
+    def __mul__(self, other: Union[PauliDesignator, PauliCoefficientDesignator]) -> 'PauliSum':
         """
         Multiplies together this PauliSum with PauliSum, PauliTerm or Number objects. The new term
         is then simplified according to the Pauli Algebra rules.
@@ -574,7 +576,7 @@ def __mul__(self, other):
         new_sum = PauliSum(new_terms)
         return new_sum.simplify()
 
-    def __rmul__(self, other):
+    def __rmul__(self, other: Number) -> 'PauliSum':
         """
         Multiples together this PauliSum with PauliSum, PauliTerm or Number objects. The new term
         is then simplified according to the Pauli Algebra rules.
@@ -589,7 +591,7 @@ def __rmul__(self, other):
             term.coefficient *= other
         return PauliSum(new_terms).simplify()
 
-    def __pow__(self, power):
+    def __pow__(self, power: int) -> 'PauliSum':
         """Raises this PauliSum to power.
 
         :param int power: The power to raise this PauliSum to.
@@ -614,7 +616,7 @@ def __pow__(self, power):
             result *= self
         return result
 
-    def __add__(self, other):
+    def __add__(self, other: Union[PauliDesignator, PauliCoefficientDesignator]) -> 'PauliSum':
         """
         Adds together this PauliSum with PauliSum, PauliTerm or Number objects. The new term
         is then simplified according to the Pauli Algebra rules.
@@ -632,7 +634,7 @@ def __add__(self, other):
         new_sum = PauliSum(new_terms)
         return new_sum.simplify()
 
-    def __radd__(self, other):
+    def __radd__(self, other: Union[PauliDesignator, PauliCoefficientDesignator]) -> 'PauliSum':
         """
         Adds together this PauliSum with PauliSum, PauliTerm or Number objects. The new term
         is then simplified according to the Pauli Algebra rules.
@@ -644,7 +646,7 @@ def __radd__(self, other):
         assert isinstance(other, Number)
         return self + other
 
-    def __sub__(self, other):
+    def __sub__(self, other: Union[PauliDesignator, PauliCoefficientDesignator]) -> 'PauliSum':
         """
         Finds the difference of this PauliSum with PauliSum, PauliTerm or Number objects. The new
         term is then simplified according to the Pauli Algebra rules.
@@ -655,7 +657,7 @@ def __sub__(self, other):
         """
         return self + -1. * other
 
-    def __rsub__(self, other):
+    def __rsub__(self, other: Union[PauliDesignator, PauliCoefficientDesignator]) -> 'PauliSum':
         """
         Finds the different of this PauliSum with PauliSum, PauliTerm or Number objects. The new
         term is then simplified according to the Pauli Algebra rules.
@@ -666,7 +668,7 @@ def __rsub__(self, other):
         """
         return other + -1. * self
 
-    def get_qubits(self):
+    def get_qubits(self) -> List[int]:
         """
         The support of all the operators in the PauliSum object.
 
@@ -675,13 +677,13 @@ def get_qubits(self):
         """
         return list(set().union(*[term.get_qubits() for term in self.terms]))
 
-    def simplify(self):
+    def simplify(self) -> 'PauliSum':
         """
         Simplifies the sum of Pauli operators according to Pauli algebra rules.
         """
         return simplify_pauli_sum(self)
 
-    def get_programs(self):
+    def get_programs(self) -> Tuple[List[Program], np.ndarray]:
         """
         Get a Pyquil Program corresponding to each term in the PauliSum and a coefficient
         for each program
@@ -692,7 +694,7 @@ def get_programs(self):
         coefficients = np.array([term.coefficient for term in self.terms])
         return programs, coefficients
 
-    def compact_str(self):
+    def compact_str(self) -> str:
         """A string representation of the PauliSum that is more compact than ``str(pauli_sum)``
 
         >>> pauli_sum = 2.0 * sX(1)* sZ(2) + 1.5 * sY(2)
@@ -704,7 +706,7 @@ def compact_str(self):
         return "+".join([term.compact_str() for term in self.terms])
 
     @classmethod
-    def from_compact_str(cls, str_pauli_sum):
+    def from_compact_str(cls, str_pauli_sum: str) -> 'PauliSum':
         """Construct a PauliSum from the result of str(pauli_sum)
         """
         # split str_pauli_sum only at "+" outside of parenthesis to allow
@@ -715,7 +717,7 @@ def from_compact_str(cls, str_pauli_sum):
         return cls(terms).simplify()
 
 
-def simplify_pauli_sum(pauli_sum):
+def simplify_pauli_sum(pauli_sum: PauliSum) -> PauliSum:
     """Simplify the sum of Pauli operators according to Pauli algebra rules."""
 
     # You might want to use a defaultdict(list) here, but don't because
@@ -748,7 +750,7 @@ def simplify_pauli_sum(pauli_sum):
     return PauliSum(terms)
 
 
-def check_commutation(pauli_list, pauli_two):
+def check_commutation(pauli_list: Sequence[PauliTerm], pauli_two: PauliTerm) -> bool:
     """
     Check if commuting a PauliTerm commutes with a list of other terms by natural calculation.
     Uses the result in Section 3 of arXiv:1405.5749v2, modified slightly here to check for the
@@ -761,7 +763,7 @@ def check_commutation(pauli_list, pauli_two):
     :rtype: bool
     """
 
-    def coincident_parity(p1, p2):
+    def coincident_parity(p1: PauliTerm, p2: PauliTerm) -> bool:
         non_similar = 0
         p1_indices = set(p1._ops.keys())
         p2_indices = set(p2._ops.keys())
@@ -776,7 +778,7 @@ def coincident_parity(p1, p2):
     return True
 
 
-def commuting_sets(pauli_terms):
+def commuting_sets(pauli_terms: PauliSum) -> List[List[PauliTerm]]:
     """Gather the Pauli terms of pauli_terms variable into commuting sets
 
     Uses algorithm defined in (Raeisi, Wiebe, Sanders, arXiv:1108.4318, 2011)
@@ -805,7 +807,7 @@ def commuting_sets(pauli_terms):
     return groups
 
 
-def is_identity(term):
+def is_identity(term: PauliDesignator) -> bool:
     """
     Tests to see if a PauliTerm or PauliSum is a scalar multiple of identity
 
@@ -822,7 +824,7 @@ def is_identity(term):
         raise TypeError("is_identity only checks PauliTerms and PauliSum objects!")
 
 
-def exponentiate(term: PauliTerm):
+def exponentiate(term: PauliTerm) -> Program:
     """
     Creates a pyQuil program that simulates the unitary evolution exp(-1j * term)
 
@@ -833,7 +835,7 @@ def exponentiate(term: PauliTerm):
     return exponential_map(term)(1.0)
 
 
-def exponential_map(term):
+def exponential_map(term: PauliTerm) -> Callable[[float], Program]:
     """
     Returns a function f(alpha) that constructs the Program corresponding to exp(-1j*alpha*term).
 
@@ -847,7 +849,7 @@ def exponential_map(term):
     coeff = term.coefficient.real
     term.coefficient = term.coefficient.real
 
-    def exp_wrap(param):
+    def exp_wrap(param: float) -> Program:
         prog = Program()
         if is_identity(term):
             prog.inst(X(0))
@@ -863,7 +865,7 @@ def exp_wrap(param):
     return exp_wrap
 
 
-def exponentiate_commuting_pauli_sum(pauli_sum):
+def exponentiate_commuting_pauli_sum(pauli_sum: PauliSum) -> Callable[[float], Program]:
     """
     Returns a function that maps all substituent PauliTerms and sums them into a program. NOTE: Use
     this function with care. Substituent PauliTerms should commute.
@@ -877,13 +879,13 @@ def exponentiate_commuting_pauli_sum(pauli_sum):
 
     fns = [exponential_map(term) for term in pauli_sum]
 
-    def combined_exp_wrap(param):
+    def combined_exp_wrap(param: float) -> Program:
         return Program([f(param) for f in fns])
 
     return combined_exp_wrap
 
 
-def _exponentiate_general_case(pauli_term, param):
+def _exponentiate_general_case(pauli_term: PauliTerm, param: float) -> Program:
     """
     Returns a Quil (Program()) object corresponding to the exponential of
     the pauli_term object, i.e. exp[-1.0j * param * pauli_term]
@@ -894,7 +896,7 @@ def _exponentiate_general_case(pauli_term, param):
     :rtype: Program
     """
 
-    def reverse_hack(p):
+    def reverse_hack(p: Program) -> Program:
         # A hack to produce a *temporary* program which reverses p.
         revp = Program()
         revp.inst(list(reversed(p.instructions)))
@@ -935,7 +937,7 @@ def reverse_hack(p):
     return quil_prog
 
 
-def suzuki_trotter(trotter_order, trotter_steps):
+def suzuki_trotter(trotter_order: int, trotter_steps: int) -> List[Tuple[float, int]]:
     """
     Generate trotterization coefficients for a given number of Trotter steps.
 
@@ -970,7 +972,7 @@ def suzuki_trotter(trotter_order, trotter_steps):
     return order_slices
 
 
-def is_zero(pauli_object):
+def is_zero(pauli_object: PauliDesignator) -> bool:
     """
     Tests to see if a PauliTerm or PauliSum is zero.
 
@@ -986,8 +988,8 @@ def is_zero(pauli_object):
         raise TypeError("is_zero only checks PauliTerms and PauliSum objects!")
 
 
-def trotterize(first_pauli_term, second_pauli_term, trotter_order=1,
-               trotter_steps=1):
+def trotterize(first_pauli_term: PauliTerm, second_pauli_term: PauliTerm, trotter_order: int = 1,
+               trotter_steps: int = 1) -> Program:
     """
     Create a Quil program that approximates exp( (A + B)t) where A and B are
     PauliTerm operators.