Extract the projected density of states from a CompleteDos object.

    Args:
        dos (:obj:`~pymatgen.electronic_structure.dos.CompleteDos`): The
            density of states.
        elements (:obj:`dict`, optional): The elements and orbitals to extract
            from the projected density of states. Should be provided as a
            :obj:`dict` with the keys as the element names and corresponding
            values as a :obj:`tuple` of orbitals. For example, the following
            would extract the Bi s, px, py and d orbitals::

                {'Bi': ('s', 'px', 'py', 'd')}

            If an element is included with an empty :obj:`tuple`, all orbitals
            for that species will be extracted. If ``elements`` is not set or
            set to ``None``, all elements for all species will be extracted.
        lm_orbitals (:obj:`dict`, optional): The orbitals to decompose into
            their lm contributions (e.g. p -> px, py, pz). Should be provided
            as a :obj:`dict`, with the elements names as keys and a
            :obj:`tuple` of orbitals as the corresponding values. For example,
            the following would be used to decompose the oxygen p and d
            orbitals::

                {'O': ('p', 'd')}

        atoms (:obj:`dict`, optional): Which atomic sites to use when
            calculating the projected density of states. Should be provided as
            a :obj:`dict`, with the element names as keys and a :obj:`tuple` of
            :obj:`int` specifying the atomic indices as the corresponding
            values. The elemental projected density of states will be summed
            only over the atom indices specified. If an element is included
            with an empty :obj:`tuple`, then all sites for that element will
            be included. The indices are 0 based for each element specified in
            the POSCAR. For example, the following will calculate the density
            of states for the first 4 Sn atoms and all O atoms in the
            structure::

                {'Sn': (1, 2, 3, 4), 'O': (, )}

            If ``atoms`` is not set or set to ``None`` then all atomic sites
            for all elements will be considered.

    Returns:
        dict: The projected density of states. Formatted as a :obj:`dict` of
        :obj:`dict` mapping the elements and their orbitals to
        :obj:`~pymatgen.electronic_structure.dos.Dos` objects. For example::

            {
                'Bi': {'s': Dos, 'p': Dos ... },
                'S': {'s': Dos}
            }

def get_pdos(dos, lm_orbitals=None, atoms=None, elements=None):
    """Extract the projected density of states from a CompleteDos object.

    Args:
        dos (:obj:`~pymatgen.electronic_structure.dos.CompleteDos`): The
            density of states.
        elements (:obj:`dict`, optional): The elements and orbitals to extract
            from the projected density of states. Should be provided as a
            :obj:`dict` with the keys as the element names and corresponding
            values as a :obj:`tuple` of orbitals. For example, the following
            would extract the Bi s, px, py and d orbitals::

                {'Bi': ('s', 'px', 'py', 'd')}

            If an element is included with an empty :obj:`tuple`, all orbitals
            for that species will be extracted. If ``elements`` is not set or
            set to ``None``, all elements for all species will be extracted.
        lm_orbitals (:obj:`dict`, optional): The orbitals to decompose into
            their lm contributions (e.g. p -> px, py, pz). Should be provided
            as a :obj:`dict`, with the elements names as keys and a
            :obj:`tuple` of orbitals as the corresponding values. For example,
            the following would be used to decompose the oxygen p and d
            orbitals::

                {'O': ('p', 'd')}

        atoms (:obj:`dict`, optional): Which atomic sites to use when
            calculating the projected density of states. Should be provided as
            a :obj:`dict`, with the element names as keys and a :obj:`tuple` of
            :obj:`int` specifying the atomic indices as the corresponding
            values. The elemental projected density of states will be summed
            only over the atom indices specified. If an element is included
            with an empty :obj:`tuple`, then all sites for that element will
            be included. The indices are 0 based for each element specified in
            the POSCAR. For example, the following will calculate the density
            of states for the first 4 Sn atoms and all O atoms in the
            structure::

                {'Sn': (1, 2, 3, 4), 'O': (, )}

            If ``atoms`` is not set or set to ``None`` then all atomic sites
            for all elements will be considered.

    Returns:
        dict: The projected density of states. Formatted as a :obj:`dict` of
        :obj:`dict` mapping the elements and their orbitals to
        :obj:`~pymatgen.electronic_structure.dos.Dos` objects. For example::

            {
                'Bi': {'s': Dos, 'p': Dos ... },
                'S': {'s': Dos}
            }
    """
    if not elements:
        symbols = dos.structure.symbol_set
        elements = dict(zip(symbols, [None] * len(symbols)))
    pdos = {}
    for el in elements:
        if atoms and el not in atoms:
            continue

        # select which sites to consider, if no sites were specified then
        # select all. Make a list of the sites of particular elements first
        # due to the dosplot atoms list specification (e.g. starts at 0 for
        # each element
        element_sites = [site for site in dos.structure.sites
                         if site.specie == get_el_sp(el)]
        sites = [site for i, site in enumerate(element_sites)
                 if not atoms or (el in atoms and i in atoms[el])]
        lm = lm_orbitals[el] if (lm_orbitals and el in lm_orbitals) else None
        orbitals = elements[el] if elements and el in elements else None

        pdos[el] = get_element_pdos(dos, el, sites, lm, orbitals)
    return pdos