def setUp(self):
     entrylist = list()
     weights = list()
     comp = Composition("Mn2O3")
     entry = PDEntry(comp, 49)
     entrylist.append(PourbaixEntry(entry))
     weights.append(1.0)
     comp = Ion.from_formula("MnO4[-]")
     entry = IonEntry(comp, 25)
     entrylist.append(PourbaixEntry(entry))
     weights.append(0.25)
     comp = Composition("Fe2O3")
     entry = PDEntry(comp, 50)
     entrylist.append(PourbaixEntry(entry))
     weights.append(0.5)
     comp = Ion.from_formula("Fe[2+]")
     entry = IonEntry(comp, 15)
     entrylist.append(PourbaixEntry(entry))
     weights.append(2.5)
     comp = Ion.from_formula("Fe[3+]")
     entry = IonEntry(comp, 20)
     entrylist.append(PourbaixEntry(entry))
     weights.append(1.5)
     self.weights = weights
     self.entrylist = entrylist
     self.multientry = MultiEntry(entrylist, weights)

 def setUp(self):
     self.comp = list()
     self.comp.append(Ion.from_formula("Li+"))
     self.comp.append(Ion.from_formula("MnO4-"))
     self.comp.append(Ion.from_formula("Mn++"))
     self.comp.append(Ion.from_formula("PO3-2"))
     self.comp.append(Ion.from_formula("Fe(CN)6-3"))
     self.comp.append(Ion.from_formula("Fe(CN)6----"))
     self.comp.append(Ion.from_formula("Fe2((PO4)3(CO3)5)2-3"))
     self.comp.append(Ion.from_formula("Ca[2+]"))
     self.comp.append(Ion.from_formula("NaOH(aq)"))

    def test_read_write_csv(self):
        Zn_solids = ["Zn", "ZnO", "ZnO2"]
        sol_g = [0.0, -3.338, -1.315]
        Zn_ions = ["Zn[2+]", "ZnOH[+]", "HZnO2[-]", "ZnO2[2-]", "ZnO"]
        liq_g = [-1.527, -3.415, -4.812, -4.036, -2.921]
        liq_conc = [1e-6, 1e-6, 1e-6, 1e-6, 1e-6]
        solid_entry = list()
        for sol in Zn_solids:
            comp = Composition(sol)
            delg = sol_g[Zn_solids.index(sol)]
            solid_entry.append(PourbaixEntry(PDEntry(comp, delg)))
        ion_entry = list()
        for ion in Zn_ions:
            comp_ion = Ion.from_formula(ion)
            delg = liq_g[Zn_ions.index(ion)]
            conc = liq_conc[Zn_ions.index(ion)]
            PoE = PourbaixEntry(IonEntry(comp_ion, delg))
            PoE.set_conc(conc)
            ion_entry.append(PoE)
        entries = solid_entry + ion_entry
        PourbaixEntryIO.to_csv("pourbaix_test_entries.csv", entries)

        (elements, entries) = PourbaixEntryIO.from_csv(
            "pourbaix_test_entries.csv")
        self.assertEqual(elements,
                         [Element('Zn'), Element('H'), Element('O')],
                         "Wrong elements!")
        self.assertEqual(len(entries), 8, "Wrong number of entries!")
        os.remove("pourbaix_test_entries.csv")

 def len_elts(entry):
     if "(s)" in entry:
         comp = Composition(entry[:-3])
     else:
         comp = Ion.from_formula(entry)
     return len([el for el in comp.elements if el not in
                 [Element("H"), Element("O")]])

 def test_equals(self):
     random_z = random.randint(1, 92)
     fixed_el = Element.from_Z(random_z)
     other_z = random.randint(1, 92)
     while other_z == random_z:
         other_z = random.randint(1, 92)
     comp1 = Ion(Composition({fixed_el: 1, Element.from_Z(other_z): 0}), 1)
     other_z = random.randint(1, 92)
     while other_z == random_z:
         other_z = random.randint(1, 92)
     comp2 = Ion(Composition({fixed_el: 1, Element.from_Z(other_z): 0}), 1)
     self.assertEqual(comp1, comp2,
                      "Composition equality test failed. " +
                      "%s should be equal to %s" % (comp1.formula,
                                                    comp2.formula))
     self.assertEqual(comp1.__hash__(), comp2.__hash__(),
                      "Hashcode equality test failed!")

 def setUp(self):
     # comp = Composition("Mn2O3")
     self.solentry = ComputedEntry("Mn2O3", 49)
     ion = Ion.from_formula("MnO4-")
     self.ionentry = IonEntry(ion, 25)
     self.PxIon = PourbaixEntry(self.ionentry)
     self.PxSol = PourbaixEntry(self.solentry)
     self.PxIon.concentration = 1e-4

 def test_as_dict(self):
     c = Ion.from_dict({'Mn': 1, 'O': 4, 'charge': -1})
     d = c.as_dict()
     correct_dict = {'Mn': 1.0, 'O': 4.0, 'charge': -1.0}
     self.assertEqual(d, correct_dict)
     self.assertEqual(d['charge'], correct_dict['charge'])
     correct_dict = {'Mn': 1.0, 'O': 4.0, 'charge': -1}
     d = c.to_reduced_dict
     self.assertEqual(d, correct_dict)
     self.assertEqual(d['charge'], correct_dict['charge'])

    def get_pourbaix_entries(self, chemsys):
        """
        A helper function to get all entries necessary to generate
        a pourbaix diagram from the rest interface.

        Args:
            chemsys ([str]): A list of elements comprising the chemical
                system, e.g. ['Li', 'Fe']
        """
        from pymatgen.analysis.pourbaix.entry import PourbaixEntry, IonEntry
        from pymatgen.analysis.phase_diagram import PhaseDiagram
        from pymatgen.core.ion import Ion
        from pymatgen.entries.compatibility import\
            MaterialsProjectAqueousCompatibility

        chemsys = list(set(chemsys + ['O', 'H']))
        entries = self.get_entries_in_chemsys(
            chemsys, property_data=['e_above_hull'], compatible_only=False)
        compat = MaterialsProjectAqueousCompatibility("Advanced")
        entries = compat.process_entries(entries)
        solid_pd = PhaseDiagram(entries) # Need this to get ion formation energy
        url = '/pourbaix_diagram/reference_data/' + '-'.join(chemsys)
        ion_data = self._make_request(url)

        pbx_entries = []
        for entry in entries:
            if not set(entry.composition.elements)\
                    <= {Element('H'), Element('O')}:
                pbx_entry = PourbaixEntry(entry)
                pbx_entry.g0_replace(solid_pd.get_form_energy(entry))
                pbx_entry.reduced_entry()
                pbx_entries.append(pbx_entry)

        # position the ion energies relative to most stable reference state
        for n, i_d in enumerate(ion_data):
            ion_entry = IonEntry(Ion.from_formula(i_d['Name']), i_d['Energy'])
            refs = [e for e in entries
                    if e.composition.reduced_formula == i_d['Reference Solid']]
            if not refs:
                raise ValueError("Reference solid not contained in entry list")
            stable_ref = sorted(refs, key=lambda x: x.data['e_above_hull'])[0]
            rf = stable_ref.composition.get_reduced_composition_and_factor()[1]
            solid_diff = solid_pd.get_form_energy(stable_ref)\
                         - i_d['Reference solid energy'] * rf
            elt = i_d['Major_Elements'][0]
            correction_factor = ion_entry.ion.composition[elt]\
                                / stable_ref.composition[elt]
            correction = solid_diff * correction_factor
            pbx_entries.append(PourbaixEntry(ion_entry, correction,
                                             'ion-{}'.format(n)))
        return pbx_entries

    def __init__(self, entries, comp_dict=None):
        """
        Args:
            entries:
                Entries list containing both Solids and Ions
            comp_dict:
                Dictionary of compositions
        """
        self._solid_entries = list()
        self._ion_entries = list()
        for entry in entries:
            if entry.phase_type == "Solid":
                self._solid_entries.append(entry)
            elif entry.phase_type == "Ion":
                self._ion_entries.append(entry)
            else:
                raise StandardError("Incorrect Phase type - needs to be \
                Pourbaix entry of phase type Ion/Solid")
        self._unprocessed_entries = self._solid_entries + self._ion_entries
        self._elt_comp = comp_dict
        if comp_dict:
            self._multielement = True
            self.pourbaix_elements = [key for key in comp_dict]
            w = [comp_dict[key] for key in comp_dict]
            A = []
            for comp in comp_dict:
                m = re.search(r"\[([^\[\]]+)\]|\(aq\)", comp)
                if m:
                    comp_obj = Ion.from_formula(comp)
                else:
                    comp_obj = Composition.from_formula(comp)
                Ai = []
                for elt in self.pourbaix_elements:
                    Ai.append(comp_obj[Element(elt)])
                A.append(Ai)
            A = np.array(A).T.astype(float)
            w = np.array(w)
            A /= np.dot([A[i].sum() for i in xrange(len(A))], w)
            x = np.linalg.solve(A, w)
            self._elt_comp = dict(zip(self.pourbaix_elements, x))

        else:
            self._multielement = False
            self.pourbaix_elements = [el.symbol
                                      for el in entries[0].composition.elements
                                      if el.symbol not in ["H", "O"]]
        self._make_pourbaixdiagram()

def ion_or_solid_comp_object(formula):
    """
    Returns either an ion object or composition object given
    a formula.

    Args:
        formula: String formula. Eg. of ion: NaOH(aq), Na[+];
            Eg. of solid: Fe2O3(s), Fe(s), Na2O

    Returns:
        Composition/Ion object
    """
    m = re.search(r"\[([^\[\]]+)\]|\(aq\)", formula)
    if m:
        comp_obj = Ion.from_formula(formula)
    elif re.search(r"\(s\)", formula):
        comp_obj = Composition(formula[:-3])
    else:
        comp_obj = Composition(formula)
    return comp_obj

    def from_csv(filename):
        """
        Imports PourbaixEntries from a csv.

        Args:
            filename: Filename to import from.

        Returns:
            List of Entries
        """
        with open(filename, "r") as f:
            reader = csv.reader(f, delimiter=unicode2str(","),
                                quotechar=unicode2str("\""),
                                quoting=csv.QUOTE_MINIMAL)
            entries = list()
            header_read = False
            for row in reader:
                if not header_read:
                    elements = row[1:(len(row) - 4)]
                    header_read = True
                else:
                    name = row[0]
                    energy = float(row[-4])
                    conc = float(row[-1])
                    comp = dict()
                    for ind in range(1, len(row) - 4):
                        if float(row[ind]) > 0:
                            comp[Element(elements[ind - 1])] = float(row[ind])
                    phase_type = row[-3]
                    if phase_type == "Ion":
                        PoE = PourbaixEntry(IonEntry(Ion.from_formula(name),
                                                     energy))
                        PoE.conc = conc
                        PoE.name = name
                        entries.append(PoE)
                    else:
                        entries.append(PourbaixEntry(PDEntry(Composition(comp),
                                                             energy)))
        elements = [Element(el) for el in elements]
        return elements, entries

    def test_mpr_pipeline(self):
        from pymatgen import MPRester
        mpr = MPRester()
        data = mpr.get_pourbaix_entries(["Zn"])
        pbx = PourbaixDiagram(data, filter_solids=True, conc_dict={"Zn": 1e-8})
        pbx.find_stable_entry(10, 0)

        data = mpr.get_pourbaix_entries(["Ag", "Te"])
        pbx = PourbaixDiagram(data, filter_solids=True,
                              conc_dict={"Ag": 1e-8, "Te": 1e-8})
        self.assertEqual(len(pbx.stable_entries), 30)
        test_entry = pbx.find_stable_entry(8, 2)
        self.assertAlmostEqual(test_entry.energy, 2.393900378500001)

        # Test custom ions
        entries = mpr.get_pourbaix_entries(["Sn", "C", "Na"])
        ion = IonEntry(Ion.from_formula("NaO28H80Sn12C24+"), -161.676)
        custom_ion_entry = PourbaixEntry(ion, entry_id='my_ion')
        pbx = PourbaixDiagram(entries + [custom_ion_entry], filter_solids=True,
                              comp_dict={"Na": 1, "Sn": 12, "C": 24})
        self.assertAlmostEqual(pbx.get_decomposition_energy(custom_ion_entry, 5, 2),
                               8.31082110278154)

def plot_pourbaix_diagram(metastability=0.0, ion_concentration=1e-6, fmt='pdf'):
    """
    Creates a Pourbaix diagram for the material in the cwd.

    Args:
        metastability (float): desired metastable tolerance energy
            (meV/atom). <~50 is generally a sensible range to use.
        ion_concentration (float): in mol/kg. Sensible values are
            generally between 1e-8 and 1.
        fmt (str): matplotlib format style. Check the matplotlib
            docs for options.
    """

    # Create a ComputedEntry object for the 2D material.
    composition = Structure.from_file('POSCAR').composition
    energy = Vasprun('vasprun.xml').final_energy

    cmpd = ComputedEntry(composition, energy)

    # Define the chemsys that describes the 2D compound.
    chemsys = ['O', 'H'] + [elt.symbol for elt in composition.elements
                            if elt.symbol not in ['O', 'H']]

    # Experimental ionic energies
    # See ions.yaml for ion formation energies and references.
    exp_dict = ION_DATA['ExpFormEnergy']
    ion_correction = ION_DATA['IonCorrection']

    # Pick out the ions pertaining to the 2D compound.
    ion_dict = dict()
    for elt in chemsys:
        if elt not in ['O', 'H'] and exp_dict[elt]:
            ion_dict.update(exp_dict[elt])

    elements = [Element(elt) for elt in chemsys if elt not in ['O', 'H']]

    # Add "correction" for metastability
    cmpd.correction -= float(cmpd.composition.num_atoms)\
        * float(metastability) / 1000.0

    # Calculate formation energy of the compound from its end
    # members
    form_energy = cmpd.energy
    for elt in composition.as_dict():
        form_energy -= END_MEMBERS[elt] * cmpd.composition[elt]

    # Convert the compound entry to a pourbaix entry.
    # Default concentration for solid entries = 1
    pbx_cmpd = PourbaixEntry(cmpd)
    pbx_cmpd.g0_replace(form_energy)
    pbx_cmpd.reduced_entry()

    # Add corrected ionic entries to the pourbaix diagram
    # dft corrections for experimental ionic energies:
    # Persson et.al PHYSICAL REVIEW B 85, 235438 (2012)
    pbx_ion_entries = list()

    # Get PourbaixEntry corresponding to each ion.
    # Default concentration for ionic entries = 1e-6
    # ion_energy = ion_exp_energy + ion_correction * factor
    # where factor = fraction of element el in the ionic entry
    # compared to the reference entry
    for elt in elements:
        for key in ion_dict:
            comp = Ion.from_formula(key)
            if comp.composition[elt] != 0:
                factor = comp.composition[elt]
                energy = ion_dict[key]
                pbx_entry_ion = PourbaixEntry(IonEntry(comp, energy))
                pbx_entry_ion.correction = ion_correction[elt.symbol]\
                    * factor
                pbx_entry_ion.conc = ion_concentration
                pbx_entry_ion.name = key
                pbx_ion_entries.append(pbx_entry_ion)

    # Generate and plot Pourbaix diagram
    # Each bulk solid/ion has a free energy g of the form:
    # g = g0_ref + 0.0591 * log10(conc) - nO * mu_H2O +
    # (nH - 2nO) * pH + phi * (-nH + 2nO + q)

    all_entries = [pbx_cmpd] + pbx_ion_entries

    pourbaix = PourbaixDiagram(all_entries)

    # Analysis features
    panalyzer = PourbaixAnalyzer(pourbaix)
#    instability = panalyzer.get_e_above_hull(pbx_cmpd)

    plotter = PourbaixPlotter(pourbaix)
    plot = plotter.get_pourbaix_plot(limits=[[0, 14], [-2, 2]],
                                     label_domains=True)
    fig = plot.gcf()
    ax1 = fig.gca()

    # Add coloring to highlight the stability region for the 2D
    # material, if one exists.
    stable_entries = plotter.pourbaix_plot_data(
        limits=[[0, 14], [-2, 2]])[0]

    for entry in stable_entries:
#.........这里部分代码省略.........

 def from_dict(cls, d):
     """
     Returns an IonEntry object from a dict.
     """
     return IonEntry(Ion.from_dict(d["composition"]), d["energy"])

 def from_dict(cls, d):
     """
     Returns an IonEntry object from a dict.
     """
     return IonEntry(Ion.from_dict(d["ion"]), d["energy"], d.get("name", None))

    def get_pourbaix_plot_colorfill_by_element(self, limits=None, title="",
                                                label_domains=True, element=None):
        """
        Color domains by element
        """
        from matplotlib.patches import Polygon
        import matplotlib.patheffects as PathEffects

        entry_dict_of_multientries = collections.defaultdict(list)
        plt = get_publication_quality_plot(16)
        optim_colors = ['#0000FF', '#FF0000', '#00FF00', '#FFFF00', '#FF00FF',
                         '#FF8080', '#DCDCDC', '#800000', '#FF8000']
        optim_font_color = ['#FFFFA0', '#00FFFF', '#FF00FF', '#0000FF', '#00FF00',
                            '#007F7F', '#232323', '#7FFFFF', '#007FFF']
        hatch = ['/', '\\', '|', '-', '+', 'o', '*']
        (stable, unstable) = self.pourbaix_plot_data(limits)
        num_of_overlaps = {key: 0 for key in stable.keys()}
        for entry in stable:
            if isinstance(entry, MultiEntry):
                for e in entry.entrylist:
                    if element in e.composition.elements:
                        entry_dict_of_multientries[e.name].append(entry)
                        num_of_overlaps[entry] += 1
            else:
                entry_dict_of_multientries[entry.name].append(entry)
        if limits:
            xlim = limits[0]
            ylim = limits[1]
        else:
            xlim = self._analyzer.chempot_limits[0]
            ylim = self._analyzer.chempot_limits[1]

        h_line = np.transpose([[xlim[0], -xlim[0] * PREFAC],
                               [xlim[1], -xlim[1] * PREFAC]])
        o_line = np.transpose([[xlim[0], -xlim[0] * PREFAC + 1.23],
                               [xlim[1], -xlim[1] * PREFAC + 1.23]])
        neutral_line = np.transpose([[7, ylim[0]], [7, ylim[1]]])
        V0_line = np.transpose([[xlim[0], 0], [xlim[1], 0]])

        ax = plt.gca()
        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        from pymatgen import Composition, Element
        from pymatgen.core.ion import Ion

        def len_elts(entry):
            if "(s)" in entry:
                comp = Composition(entry[:-3])
            else:
                comp = Ion.from_formula(entry)
            return len([el for el in comp.elements if el not in
                        [Element("H"), Element("O")]])

        sorted_entry = entry_dict_of_multientries.keys()
        sorted_entry.sort(key=len_elts)
        i = -1
        label_chr = map(chr, range(65, 91))
        for entry in sorted_entry:
            color_indx = 0
            x_coord = 0.0
            y_coord = 0.0
            npts = 0
            i += 1
            for e in entry_dict_of_multientries[entry]:
                hc = 0
                fc = 0
                bc = 0
                xy = self.domain_vertices(e)
                c = self.get_center(stable[e])
                x_coord += c[0]
                y_coord += c[1]
                npts += 1
                color_indx = i
                if "(s)" in entry:
                    comp = Composition(entry[:-3])
                else:
                    comp = Ion.from_formula(entry)
                if len([el for el in comp.elements if el not in
                         [Element("H"), Element("O")]]) == 1:
                    if color_indx >= len(optim_colors):
                        color_indx = color_indx -\
                         int(color_indx / len(optim_colors)) * len(optim_colors)
                    patch = Polygon(xy, facecolor=optim_colors[color_indx],
                                     closed=True, lw=3.0, fill=True)
                    bc = optim_colors[color_indx]
                else:
                    if color_indx >= len(hatch):
                        color_indx = color_indx - int(color_indx / len(hatch)) * len(hatch)
                    patch = Polygon(xy, hatch=hatch[color_indx], closed=True, lw=3.0, fill=False)
                    hc = hatch[color_indx]
                ax.add_patch(patch)
                
            xy_center = (x_coord / npts, y_coord / npts)
            if label_domains:
                if color_indx >= len(optim_colors):
                    color_indx = color_indx -\
                        int(color_indx / len(optim_colors)) * len(optim_colors)
                fc = optim_font_color[color_indx]
                if bc and not hc:
                    bbox = dict(boxstyle="round", fc=fc)
#.........这里部分代码省略.........

 def len_elts(entry):
     comp = Composition(entry[:-3]) if "(s)" in entry else Ion.from_formula(entry)
     return len(set(comp.elements) - {Element("H"), Element("O")})

    def get_pourbaix_entries(self, chemsys):
        """
        A helper function to get all entries necessary to generate
        a pourbaix diagram from the rest interface.

        Args:
            chemsys ([str]): A list of elements comprising the chemical
                system, e.g. ['Li', 'Fe']
        """
        from pymatgen.analysis.pourbaix_diagram import PourbaixEntry, IonEntry
        from pymatgen.analysis.phase_diagram import PhaseDiagram
        from pymatgen.core.ion import Ion
        from pymatgen.entries.compatibility import\
            MaterialsProjectAqueousCompatibility

        pbx_entries = []

        # Get ion entries first, because certain ions have reference
        # solids that aren't necessarily in the chemsys (Na2SO4)
        url = '/pourbaix_diagram/reference_data/' + '-'.join(chemsys)
        ion_data = self._make_request(url)
        ion_ref_comps = [Composition(d['Reference Solid']) for d in ion_data]
        ion_ref_elts = list(itertools.chain.from_iterable(
            i.elements for i in ion_ref_comps))
        ion_ref_entries = self.get_entries_in_chemsys(
            list(set([str(e) for e in ion_ref_elts] + ['O', 'H'])),
            property_data=['e_above_hull'], compatible_only=False)
        compat = MaterialsProjectAqueousCompatibility("Advanced")
        ion_ref_entries = compat.process_entries(ion_ref_entries)
        ion_ref_pd = PhaseDiagram(ion_ref_entries)

        # position the ion energies relative to most stable reference state
        for n, i_d in enumerate(ion_data):
            ion_entry = IonEntry(Ion.from_formula(i_d['Name']), i_d['Energy'])
            refs = [e for e in ion_ref_entries
                    if e.composition.reduced_formula == i_d['Reference Solid']]
            if not refs:
                raise ValueError("Reference solid not contained in entry list")
            stable_ref = sorted(refs, key=lambda x: x.data['e_above_hull'])[0]
            rf = stable_ref.composition.get_reduced_composition_and_factor()[1]
            solid_diff = ion_ref_pd.get_form_energy(stable_ref)\
                         - i_d['Reference solid energy'] * rf
            elt = i_d['Major_Elements'][0]
            correction_factor = ion_entry.ion.composition[elt]\
                                / stable_ref.composition[elt]
            ion_entry.energy += solid_diff * correction_factor
            pbx_entries.append(PourbaixEntry(ion_entry, 'ion-{}'.format(n)))
            # import nose; nose.tools.set_trace()

        # Construct the solid pourbaix entries from filtered ion_ref entries
        extra_elts = set(ion_ref_elts) - {Element(s) for s in chemsys}\
            - {Element('H'), Element('O')}
        for entry in ion_ref_entries:
            entry_elts = set(entry.composition.elements)
            # Ensure no OH chemsys or extraneous elements from ion references
            if not (entry_elts <= {Element('H'), Element('O')} or \
                    extra_elts.intersection(entry_elts)):
                # replace energy with formation energy, use dict to
                # avoid messing with the ion_ref_pd and to keep all old params
                form_e = ion_ref_pd.get_form_energy(entry)
                new_entry = deepcopy(entry)
                new_entry.uncorrected_energy = form_e
                new_entry.correction = 0.0
                pbx_entry = PourbaixEntry(new_entry)
                if entry.entry_id == "mp-697146":
                    pass
                    # import nose; nose.tools.set_trace()
                # pbx_entry.reduced_entry()
                pbx_entries.append(pbx_entry)

        return pbx_entries

 def test_from_dict(self):
     sym_dict = {"P": 1, "O": 4, 'charge': -2}
     self.assertEqual(Ion.from_dict(sym_dict).reduced_formula,
                      "PO4[2-]",
                      "Creation form sym_amount dictionary failed!")

 def test_get_composition(self):
     comp = self.entry.composition
     expected_comp = Ion.from_formula('MnO4[-]')
     self.assertEquals(comp, expected_comp, "Wrong composition!")