diff --git a/src/openfermion/chem/pubchem.py b/src/openfermion/chem/pubchem.py index ca365c431..c221bcd6e 100644 --- a/src/openfermion/chem/pubchem.py +++ b/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) diff --git a/src/openfermion/chem/pubchem_test.py b/src/openfermion/chem/pubchem_test.py index 5c3cbbffb..51017213e 100644 --- a/src/openfermion/chem/pubchem_test.py +++ b/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) diff --git a/src/openfermion/ops/operators/fermion_operator_test.py b/src/openfermion/ops/operators/fermion_operator_test.py index 13e57c465..293851986 100644 --- a/src/openfermion/ops/operators/fermion_operator_test.py +++ b/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) diff --git a/src/openfermion/ops/operators/symbolic_operator.py b/src/openfermion/ops/operators/symbolic_operator.py index fbd8d6a81..d9aa27f37 100644 --- a/src/openfermion/ops/operators/symbolic_operator.py +++ b/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): diff --git a/src/openfermion/ops/operators/symbolic_operator_test.py b/src/openfermion/ops/operators/symbolic_operator_test.py index a89cb9f21..84c36226f 100644 --- a/src/openfermion/ops/operators/symbolic_operator_test.py +++ b/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