diff --git a/mp_api/client/mprester.py b/mp_api/client/mprester.py index 1d9afc5c..952877dd 100644 --- a/mp_api/client/mprester.py +++ b/mp_api/client/mprester.py @@ -39,6 +39,7 @@ from emmet.core.tasks import CoreTaskDoc from pymatgen.analysis.phase_diagram import PDEntry + from pymatgen.entries.compatibility import Compatibility from pymatgen.entries.computed_entries import ComputedEntry @@ -611,7 +612,11 @@ def get_entries( def get_pourbaix_entries( self, chemsys: str | list[str] | list[ComputedEntry | ComputedStructureEntry], - solid_compat="MaterialsProject2020Compatibility", + solid_compat: Literal[ + "MaterialsProjectCompatibility", "MaterialsProject2020Compatibility" + ] + | Compatibility + | None = "MaterialsProject2020Compatibility", use_gibbs: Literal[300] | None = None, ): """A helper function to get all entries necessary to generate @@ -627,10 +632,12 @@ def get_pourbaix_entries( for adding extra calculation data to the Pourbaix Diagram. If this is set, the chemsys will be inferred from the entries. solid_compat: Compatibility scheme used to pre-process solid DFT energies prior - to applying aqueous energy adjustments. May be passed as a class (e.g. - MaterialsProject2020Compatibility) or an instance - (e.g., MaterialsProject2020Compatibility()). If None, solid DFT energies - are used as-is. Default: MaterialsProject2020Compatibility + to applying aqueous energy adjustments. + May be passed as a string (either "MaterialsProjectCompatibility" + or "MaterialsProject2020Compatibility"), or as a class instance + (e.g., MaterialsProject2020Compatibility()). + If None, solid DFT energies are used as-is. + Default: MaterialsProject2020Compatibility use_gibbs: Set to 300 (for 300 Kelvin) to use a machine learning model to estimate solid free energy from DFT energy (see GibbsComputedStructureEntry). This can slightly improve the accuracy of the Pourbaix diagram in some @@ -673,12 +680,12 @@ def get_pourbaix_entries( solid_compat = MaterialsProjectCompatibility() elif solid_compat == "MaterialsProject2020Compatibility": solid_compat = MaterialsProject2020Compatibility() - elif isinstance(solid_compat, Compatibility): + elif isinstance(solid_compat, Compatibility) or solid_compat is None: pass else: raise ValueError( "Solid compatibility can only be 'MaterialsProjectCompatibility', " - "'MaterialsProject2020Compatibility', or an instance of a Compatibility class" + "'MaterialsProject2020Compatibility', None, or an instance of a Compatibility class" ) pbx_entries = [] diff --git a/tests/client/test_mprester.py b/tests/client/test_mprester.py index 63774a49..6b37b987 100644 --- a/tests/client/test_mprester.py +++ b/tests/client/test_mprester.py @@ -270,7 +270,7 @@ def test_get_pourbaix_entries(self, mpr): # test solid_compat kwarg with pytest.raises(ValueError, match="Solid compatibility can only be"): - mpr.get_pourbaix_entries("Ti-O", solid_compat=None) + mpr.get_pourbaix_entries("Ti-O", solid_compat="None") # test removal of extra elements from reference solids # Li-Zn-S has Na in reference solids