Fix #481: add retry logic to pubchem fetching function

src/openfermion/chem/pubchem.py

@@ -9,6 +9,22 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+import time
+from pubchempy import get_compounds, PubChemHTTPError, ServerError
+
+
+def _get_compounds_with_retry(name: str, record_type: str):
+    """Get compounds from PubChem with retry logic."""
+    max_retries = 3
+    retry_delay = 5  # seconds
+    for i in range(max_retries):
+        try:
+            return get_compounds(name, 'name', record_type=record_type)
+        except (PubChemHTTPError, ConnectionError, ServerError) as e:
+            if i < max_retries - 1:
+                time.sleep(retry_delay)
+            else:
+                raise e
 
 
 def geometry_from_pubchem(name: str, structure: str = None):
@@ -30,17 +46,15 @@ def geometry_from_pubchem(name: str, structure: str = None):
         geometry: a list of tuples giving the coordinates of each atom with
         distances in Angstrom.
     """
-    import pubchempy
-
     if structure in ['2d', '3d']:
-        pubchempy_molecule = pubchempy.get_compounds(name, 'name', record_type=structure)
+        pubchempy_molecule = _get_compounds_with_retry(name, record_type=structure)
     elif structure is None:
         # Ideally get the 3-D geometry if available.
-        pubchempy_molecule = pubchempy.get_compounds(name, 'name', record_type='3d')
+        pubchempy_molecule = _get_compounds_with_retry(name, record_type='3d')
 
         # If the 3-D geometry isn't available, get the 2-D geometry instead.
         if not pubchempy_molecule:
-            pubchempy_molecule = pubchempy.get_compounds(name, 'name', record_type='2d')
+            pubchempy_molecule = _get_compounds_with_retry(name, record_type='2d')
     else:
         raise ValueError('Incorrect value for the argument structure=%s' % structure)
 

src/openfermion/chem/pubchem_test.py

@@ -8,14 +8,28 @@
 #   distributed under the License is distributed on an "AS IS" BASIS,
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY, either express or implied.
+#   See the License for the specific language governing permissions and
 #   limitations under the License.
 """Tests for pubchem.py."""
 
+import json
 import unittest
+from unittest.mock import Mock, patch
+
 import numpy
 import pytest
+from pubchempy import ServerError
 
-from openfermion.chem.pubchem import geometry_from_pubchem
+from openfermion.chem.pubchem import _get_compounds_with_retry, geometry_from_pubchem
 from openfermion.testing.testing_utils import module_importable
 
 using_pubchempy = pytest.mark.skipif(
@@ -25,7 +39,22 @@
 
 @using_pubchempy
 class OpenFermionPubChemTest(unittest.TestCase):
-    def test_water(self):
+    @patch('openfermion.chem.pubchem._get_compounds_with_retry')
+    def test_water(self, mock_get_compounds):
+        mock_compound = Mock()
+        mock_compound.to_dict.return_value = {
+            'atoms': [
+                {'element': 'O', 'x': 0.0, 'y': 0.0, 'z': 0.0},
+                {'element': 'H', 'x': 0.95, 'y': 0.0, 'z': 0.0},
+                {
+                    'element': 'H',
+                    'x': 0.95 * numpy.cos(104.5 * numpy.pi / 180.0),
+                    'y': 0.95 * numpy.sin(104.5 * numpy.pi / 180.0),
+                    'z': 0.0,
+                },
+            ]
+        }
+        mock_get_compounds.return_value = [mock_compound]
         water_geometry = geometry_from_pubchem('water')
         self.water_natoms = len(water_geometry)
         self.water_atoms = [water_atom[0] for water_atom in water_geometry]
@@ -40,14 +69,8 @@ def test_water(self):
         self.water_bond_length_1 = numpy.linalg.norm(water_oxygen_hydrogen1)
         self.water_bond_length_2 = numpy.linalg.norm(water_oxygen_hydrogen2)
         self.water_bond_angle = numpy.arccos(
-            numpy.dot(
-                water_oxygen_hydrogen1,
-                water_oxygen_hydrogen2
-                / (
-                    numpy.linalg.norm(water_oxygen_hydrogen1)
-                    * numpy.linalg.norm(water_oxygen_hydrogen2)
-                ),
-            )
+            numpy.dot(water_oxygen_hydrogen1, water_oxygen_hydrogen2)
+            / (self.water_bond_length_1 * self.water_bond_length_2)
         )
 
         water_natoms = 3
@@ -64,19 +87,35 @@ def test_water(self):
         self.assertTrue(water_bond_angle_low <= self.water_bond_angle)
         self.assertTrue(water_bond_angle_high >= self.water_bond_angle)
 
-    def test_helium(self):
+    @patch('openfermion.chem.pubchem._get_compounds_with_retry')
+    def test_helium(self, mock_get_compounds):
+        mock_compound = Mock()
+        mock_compound.to_dict.return_value = {'atoms': [{'element': 'He', 'x': 0, 'y': 0, 'z': 0}]}
+        mock_get_compounds.return_value = [mock_compound]
         helium_geometry = geometry_from_pubchem('helium')
         self.helium_natoms = len(helium_geometry)
 
         helium_natoms = 1
         self.assertEqual(helium_natoms, self.helium_natoms)
 
-    def test_none(self):
+    @patch('openfermion.chem.pubchem._get_compounds_with_retry')
+    def test_none(self, mock_get_compounds):
+        mock_get_compounds.return_value = None
         none_geometry = geometry_from_pubchem('none')
 
         self.assertIsNone(none_geometry)
 
-    def test_water_2d(self):
+    @patch('openfermion.chem.pubchem._get_compounds_with_retry')
+    def test_water_2d(self, mock_get_compounds):
+        mock_compound = Mock()
+        mock_compound.to_dict.return_value = {
+            'atoms': [
+                {'element': 'O', 'x': 0, 'y': 0, 'z': 0},
+                {'element': 'H', 'x': 1, 'y': 0, 'z': 0},
+                {'element': 'H', 'x': 0, 'y': 1, 'z': 0},
+            ]
+        }
+        mock_get_compounds.return_value = [mock_compound]
         water_geometry = geometry_from_pubchem('water', structure='2d')
         self.water_natoms = len(water_geometry)
 
@@ -91,3 +130,17 @@ def test_water_2d(self):
 
         with pytest.raises(ValueError, match='Incorrect value for the argument structure'):
             _ = geometry_from_pubchem('water', structure='foo')
+
+    @patch('openfermion.chem.pubchem.get_compounds')
+    def test_retry_logic(self, mock_get_compounds):
+        mock_get_compounds.side_effect = [ServerError('Error'), 'Success']
+        result = _get_compounds_with_retry('water', '3d')
+        self.assertEqual(result, 'Success')
+        self.assertEqual(mock_get_compounds.call_count, 2)
+
+    @patch('openfermion.chem.pubchem.get_compounds')
+    def test_retry_logic_exception(self, mock_get_compounds):
+        mock_get_compounds.side_effect = ServerError('Error')
+        with pytest.raises(ServerError):
+            _get_compounds_with_retry('water', '3d')
+        self.assertEqual(mock_get_compounds.call_count, 3)

src/openfermion/ops/operators/fermion_operator_test.py

@@ -9,11 +9,15 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-"""Tests  fermion_operator.py."""
+
+"""Tests fermion_operator.py."""
+
 import unittest
 
-from openfermion.ops.operators.fermion_operator import FermionOperator
+import sympy
+
 from openfermion.hamiltonians import number_operator
+from openfermion.ops.operators.fermion_operator import FermionOperator
 
 
 class FermionOperatorTest(unittest.TestCase):
@@ -78,3 +82,10 @@ def test_is_two_body_number_conserving_three(self):
     def test_is_two_body_number_conserving_out_of_order(self):
         op = FermionOperator(((0, 1), (2, 0), (1, 1), (3, 0)))
         self.assertTrue(op.is_two_body_number_conserving())
+
+    def test_add_sympy_rational(self):
+        """Test adding operators with sympy.Rational coefficients."""
+        a = FermionOperator('0^ 0', sympy.Rational(1, 2))
+        b = FermionOperator('1^ 1', sympy.Rational(1, 2))
+        c = a + b
+        self.assertIsInstance(c.terms[((0, 1), (0, 0))], sympy.Rational)

src/openfermion/ops/operators/symbolic_operator.py

@@ -433,7 +433,7 @@ def __iadd__(self, addend):
         """
         if isinstance(addend, type(self)):
             for term in addend.terms:
-                self.terms[term] = self.terms.get(term, 0.0) + addend.terms[term]
+                self.terms[term] = self.terms.get(term, 0) + addend.terms[term]
                 if self._issmall(self.terms[term]):
                     del self.terms[term]
         elif isinstance(addend, COEFFICIENT_TYPES):
@@ -480,7 +480,7 @@ def __isub__(self, subtrahend):
         """
         if isinstance(subtrahend, type(self)):
             for term in subtrahend.terms:
-                self.terms[term] = self.terms.get(term, 0.0) - subtrahend.terms[term]
+                self.terms[term] = self.terms.get(term, 0) - subtrahend.terms[term]
                 if self._issmall(self.terms[term]):
                     del self.terms[term]
         elif isinstance(subtrahend, COEFFICIENT_TYPES):

src/openfermion/ops/operators/symbolic_operator_test.py

@@ -9,17 +9,20 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+
 """Tests  symbolic_operator.py."""
+
 import copy
 import unittest
 import warnings
 
 import numpy
 import sympy
-from openfermion.config import EQ_TOLERANCE
-from openfermion.testing.testing_utils import EqualsTester
 
+from openfermion.config import EQ_TOLERANCE
+from openfermion.ops.operators.fermion_operator import FermionOperator
 from openfermion.ops.operators.symbolic_operator import SymbolicOperator
+from openfermion.testing.testing_utils import EqualsTester
 
 
 class DummyOperator1(SymbolicOperator):
@@ -687,6 +690,14 @@ def test_add_sympy(self):
         self.assertTrue(a.terms[term_a] - coeff_a == 0)
         self.assertTrue(a.terms[term_b] - coeff_b - 0.5 == 0)
 
+    def test_add_sympy_new_term(self):
+        """Test adding a new term with a sympy coefficient."""
+        x = sympy.Symbol('x')
+        op = FermionOperator('1^', x)
+        op += FermionOperator('2', 2 * x)
+        self.assertEqual(op.terms[((1, 1),)], x)
+        self.assertEqual(op.terms[((2, 0),)], 2 * x)
+
     def test_radd(self):
         term_a = ((1, 1), (3, 0), (8, 1))
         coeff_a = 1