Accept NumPy scalar coefficients across operator and representation types

src/openfermion/config.py

@@ -10,11 +10,17 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import numbers
 import os
 
 # Tolerance to consider number zero.
 EQ_TOLERANCE = 1e-8
 
+# Numeric types accepted as operator and tensor coefficients. numbers.Number
+# covers Python and NumPy scalar types alike (NumPy scalars are not subclasses
+# of the built-in int/float/complex types).
+COEFFICIENT_TYPES = (int, float, complex, numbers.Number)
+
 # Molecular data directory.
 THIS_DIRECTORY = os.path.dirname(os.path.realpath(__file__))
 DATA_DIRECTORY = os.path.realpath(os.path.join(THIS_DIRECTORY, 'testing/data'))

src/openfermion/hamiltonians/special_operators.py

@@ -13,9 +13,12 @@
 
 from typing import Optional, Union, Tuple
 
+import openfermion.config as config
 from openfermion.ops.operators import BosonOperator, FermionOperator
 from openfermion.utils.indexing import down_index, up_index
 
+COEFFICIENT_TYPES = config.COEFFICIENT_TYPES
+
 
 def s_plus_operator(n_spatial_orbitals: int) -> FermionOperator:
     r"""Return the s+ operator.
@@ -236,7 +239,7 @@ def majorana_operator(
     Returns:
         FermionOperator
     """
-    if not isinstance(coefficient, (int, float, complex)):
+    if not isinstance(coefficient, COEFFICIENT_TYPES):
         raise ValueError('Coefficient must be scalar.')
 
     # If term is a string, convert it to a tuple

src/openfermion/hamiltonians/special_operators_test.py

@@ -12,6 +12,9 @@
 """testing angular momentum generators. _fermion_spin_operators.py"""
 
 import unittest
+
+import numpy
+
 from openfermion.ops.operators import FermionOperator, BosonOperator
 from openfermion.utils import commutator
 from openfermion.transforms.opconversions import normal_ordered
@@ -204,3 +207,11 @@ def test_bad_term(self):
             majorana_operator('a')
         with self.assertRaises(ValueError):
             majorana_operator(2)
+
+    def test_builder_numpy_scalar_coefficient(self):
+        """The majorana_operator builder accepts NumPy scalar coefficients (issue #1097)."""
+        cases = [(numpy.int64(2), 2), (numpy.float32(0.5), 0.5), (numpy.complex64(1 + 2j), 1 + 2j)]
+        for numpy_scalar, python_scalar in cases:
+            self.assertEqual(
+                majorana_operator((1, 0), numpy_scalar), majorana_operator((1, 0), python_scalar)
+            )

src/openfermion/ops/operators/majorana_operator.py

@@ -13,8 +13,13 @@
 
 import copy
 import itertools
+
 import numpy
 
+import openfermion.config as config
+
+COEFFICIENT_TYPES = config.COEFFICIENT_TYPES
+
 
 class MajoranaOperator:
     r"""A linear combination of products of Majorana operators.
@@ -78,7 +83,7 @@ def from_dict(terms):
 
     def commutes_with(self, other):
         """Test commutation with another MajoranaOperator"""
-        if isinstance(other, (int, float, complex)):
+        if isinstance(other, COEFFICIENT_TYPES):
             return True
 
         if not isinstance(other, type(self)):
@@ -145,7 +150,7 @@ def __iadd__(self, other):
                     self.terms[term] += coefficient
                 else:
                     self.terms[term] = coefficient
-        elif isinstance(other, (int, float, complex)):
+        elif isinstance(other, COEFFICIENT_TYPES):
             self.constant += other
         else:
             raise TypeError("Cannot add invalid type to {}".format(type(self)))
@@ -163,7 +168,7 @@ def __isub__(self, other):
                     self.terms[term] -= coefficient
                 else:
                     self.terms[term] = -coefficient
-        elif isinstance(other, (int, float, complex)):
+        elif isinstance(other, COEFFICIENT_TYPES):
             self.constant -= other
         else:
             raise TypeError("Cannot subtract invalid type from {}".format(type(self)))
@@ -175,10 +180,10 @@ def __sub__(self, other):
         return minuend
 
     def __mul__(self, other):
-        if not isinstance(other, (type(self), int, float, complex)):
+        if not isinstance(other, _MAJORANA_MUL_TYPES):
             return NotImplemented
 
-        if isinstance(other, (int, float, complex)):
+        if isinstance(other, COEFFICIENT_TYPES):
             terms = {term: coefficient * other for term, coefficient in self.terms.items()}
             return MajoranaOperator.from_dict(terms)
 
@@ -194,30 +199,30 @@ def __mul__(self, other):
         return MajoranaOperator.from_dict(terms)
 
     def __imul__(self, other):
-        if not isinstance(other, (type(self), int, float, complex)):
+        if not isinstance(other, _MAJORANA_MUL_TYPES):
             return NotImplemented
 
-        if isinstance(other, (int, float, complex)):
+        if isinstance(other, COEFFICIENT_TYPES):
             for term in self.terms:
                 self.terms[term] *= other
             return self
 
         return self * other
 
     def __rmul__(self, other):
-        if not isinstance(other, (int, float, complex)):
+        if not isinstance(other, COEFFICIENT_TYPES):
             return NotImplemented
         return self * other
 
     def __truediv__(self, other):
-        if not isinstance(other, (int, float, complex)):
+        if not isinstance(other, COEFFICIENT_TYPES):
             return NotImplemented
 
         terms = {term: coefficient / other for term, coefficient in self.terms.items()}
         return MajoranaOperator.from_dict(terms)
 
     def __itruediv__(self, other):
-        if not isinstance(other, (int, float, complex)):
+        if not isinstance(other, COEFFICIENT_TYPES):
             return NotImplemented
 
         for term in self.terms:
@@ -275,6 +280,12 @@ def __repr__(self):
         return 'MajoranaOperator.from_dict(terms={!r})'.format(self.terms)
 
 
+# Types accepted by MajoranaOperator multiplication: another MajoranaOperator or
+# a scalar coefficient. Defined here, after the class, so the tuple is built once
+# at import rather than on every __mul__/__imul__ call.
+_MAJORANA_MUL_TYPES = (MajoranaOperator,) + COEFFICIENT_TYPES
+
+
 def _sort_majorana_term(term):
     """Sort a Majorana term.
 

src/openfermion/ops/operators/majorana_operator_test.py

@@ -238,3 +238,27 @@ def test_majorana_operator_str():
 def test_majorana_operator_repr():
     a = MajoranaOperator((0, 1, 5), 1.5)
     assert repr(a) == 'MajoranaOperator.from_dict(terms={(0, 1, 5): 1.5})'
+
+
+def test_majorana_operator_numpy_scalar_coefficients():
+    """NumPy scalar coefficients behave like Python scalars (issue #1097)."""
+    op = MajoranaOperator((0, 1), 1.0) + MajoranaOperator((2, 3), 2.0)
+    cases = [(numpy.int64(2), 2), (numpy.float32(0.5), 0.5), (numpy.complex64(1 + 2j), 1 + 2j)]
+    for numpy_scalar, python_scalar in cases:
+        assert op * numpy_scalar == op * python_scalar
+        assert numpy_scalar * op == python_scalar * op
+        assert op + numpy_scalar == op + python_scalar
+        assert op - numpy_scalar == op - python_scalar
+        assert op / numpy_scalar == op / python_scalar
+    assert op.commutes_with(numpy.int64(5))
+
+    # In-place operators, using exact values so the comparison is not subject
+    # to float32 round-off across the sequence of operations.
+    for numpy_scalar, python_scalar in [(numpy.int64(2), 2), (numpy.float32(0.5), 0.5)]:
+        numpy_op = MajoranaOperator((0, 1), 1.0) + MajoranaOperator((2, 3), 2.0)
+        python_op = MajoranaOperator((0, 1), 1.0) + MajoranaOperator((2, 3), 2.0)
+        numpy_op *= numpy_scalar
+        python_op *= python_scalar
+        numpy_op /= numpy_scalar
+        python_op /= python_scalar
+        assert numpy_op == python_op

src/openfermion/ops/representations/doci_hamiltonian.py

@@ -13,10 +13,11 @@
 
 import numpy
 
+import openfermion.config as config
 from openfermion.ops import QubitOperator
 from openfermion.ops.representations import PolynomialTensor, get_tensors_from_integrals
 
-COEFFICIENT_TYPES = (int, float, complex)
+COEFFICIENT_TYPES = config.COEFFICIENT_TYPES
 
 
 class DOCIHamiltonian(PolynomialTensor):

src/openfermion/ops/representations/doci_hamiltonian_test.py

@@ -277,3 +277,18 @@ def test_from_integrals_to_qubit(self):
             + "Hamiltonian\n"
             + str(sub_matrix)
         )
+
+    def test_numpy_scalar_coefficients(self):
+        """NumPy scalar coefficients behave like Python scalars (issue #1097)."""
+        doci = DOCIHamiltonian(
+            1.0,
+            numpy.array([1.0, 2.0]),
+            numpy.array([[0.0, 0.5], [0.5, 0.0]]),
+            numpy.array([[0.0, 0.3], [0.3, 0.0]]),
+        )
+        cases = [(numpy.int64(2), 2), (numpy.float32(0.5), 0.5)]
+        for numpy_scalar, python_scalar in cases:
+            self.assertEqual(doci * numpy_scalar, doci * python_scalar)
+            self.assertEqual(doci + numpy_scalar, doci + python_scalar)
+            self.assertEqual(doci - numpy_scalar, doci - python_scalar)
+            self.assertEqual(doci / numpy_scalar, doci / python_scalar)

src/openfermion/ops/representations/polynomial_tensor.py

@@ -18,9 +18,10 @@
 
 import numpy
 
-from openfermion.config import EQ_TOLERANCE
+import openfermion.config as config
 
-COEFFICIENT_TYPES = (int, float, complex)
+EQ_TOLERANCE = config.EQ_TOLERANCE
+COEFFICIENT_TYPES = config.COEFFICIENT_TYPES
 
 
 class PolynomialTensorError(Exception):

src/openfermion/ops/representations/polynomial_tensor_test.py

@@ -536,3 +536,19 @@ def do_rotate_basis_high_order(self, order):
         polynomial_tensor.rotate_basis(numpy.array([[rotation]]))
 
         return polynomial_tensor, want_polynomial_tensor
+
+    def test_numpy_scalar_coefficients(self):
+        """NumPy scalar coefficients behave like Python scalars (issue #1097)."""
+        tensor = PolynomialTensor(
+            {
+                (): 1.0,
+                (1, 0): numpy.array([[1.0, 2.0], [3.0, 4.0]]),
+                (1, 1, 0, 0): numpy.arange(16, dtype=float).reshape((2, 2, 2, 2)),
+            }
+        )
+        cases = [(numpy.int64(2), 2), (numpy.int32(3), 3), (numpy.float32(0.5), 0.5)]
+        for numpy_scalar, python_scalar in cases:
+            self.assertEqual(tensor * numpy_scalar, tensor * python_scalar)
+            self.assertEqual(tensor + numpy_scalar, tensor + python_scalar)
+            self.assertEqual(tensor - numpy_scalar, tensor - python_scalar)
+            self.assertEqual(tensor / numpy_scalar, tensor / python_scalar)