Source code for pymatgen.apps.battery.plotter

# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.


"""
This module provides plotting capabilities for battery related applications.
"""


__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyuep@gmail.com"
__date__ = "Jul 12, 2012"


from collections import OrderedDict
from pymatgen.util.plotting import pretty_plot


class VoltageProfilePlotter:
    """
    A plotter to make voltage profile plots for batteries.
    """

    def __init__(self, xaxis="capacity"):
        """
        Args:
            xaxis: The quantity to use as the xaxis. Can be either capacity (the
                default), or the frac_x.
        """
        self._electrodes = OrderedDict()
        self.xaxis = xaxis

    def add_electrode(self, electrode, label=None):
        """
        Add an electrode to the plot.

        Args:
            electrode: An electrode. All electrodes satisfying the
                AbstractElectrode interface should work.
            label: A label for the electrode. If None, defaults to a counting
                system, i.e. 'Electrode 1', 'Electrode 2', ...
        """
        if not label:
            label = "Electrode {}".format(len(self._electrodes) + 1)
        self._electrodes[label] = electrode

    def get_plot_data(self, electrode):
        """
        Args:
            electrode (): Electrode object

        Returns:
            Plot data in x, y.
        """
        x = []
        y = []
        cap = 0
        most_discharged = electrode[-1].frac_discharge
        norm = most_discharged / (1 - most_discharged)
        for vpair in electrode:
            if self.xaxis == "capacity":
                x.append(cap)
                cap += vpair.mAh / electrode.normalization_mass
                x.append(cap)
            else:
                x.append(vpair.frac_charge / (1 - vpair.frac_charge) / norm)
                x.append(vpair.frac_discharge / (1 - vpair.frac_discharge)
                         / norm)
            y.extend([vpair.voltage] * 2)

        x.append(x[-1])
        y.append(0)
        return x, y

    def get_plot(self, width=8, height=8):
        """
        Returns a plot object.

        Args:
            width: Width of the plot. Defaults to 8 in.
            height: Height of the plot. Defaults to 6 in.

        Returns:
            A matplotlib plot object.
        """
        plt = pretty_plot(width, height)
        for label, electrode in self._electrodes.items():
            (x, y) = self.get_plot_data(electrode)
            plt.plot(x, y, '-', linewidth=2, label=label)

        plt.legend()
        if self.xaxis == "capacity":
            plt.xlabel('Capacity (mAh/g)')
        else:
            plt.xlabel('Fraction')
        plt.ylabel('Voltage (V)')
        plt.tight_layout()
        return plt

    def show(self, width=8, height=6):
        """
        Show the voltage profile plot.

        Args:
            width: Width of the plot. Defaults to 8 in.
            height: Height of the plot. Defaults to 6 in.
        """
        self.get_plot(width, height).show()

    def save(self, filename, image_format="eps", width=8, height=6):
        """
        Save the plot to an image file.

        Args:
            filename: Filename to save to.
            image_format: Format to save to. Defaults to eps.
        """
        self.get_plot(width, height).savefig(filename, format=image_format)