Quickly plot colormap or a color palette with matplotlib

I present here simple convenient functions to plot a matplotlib colormap or a color palette.

The source codes of the functions are available on this gist.

For colormaps

You can use either its name or the matplotlib.cm.colormap object.

This is how to use the function :

fig = plot_cmap("summer", 6)


fig = plot_cmap(plt.cm.summer, 6)

And you will get :


The source code
import matplotlib.pyplot as plt
def plot_cmap(cmap, ncolor):
    A convenient function to plot colors of a matplotlib cmap
        ncolor (int): number of color to show
        cmap: a cmap object or a matplotlib color name
    if isinstance(cmap, str):
            cm = plt.get_cmap(cmap)
        except ValueError:
            print("WARNINGS :", cmap, " is not a known colormap")
            cm = plt.cm.gray
        cm = cmap
    with mpl.rc_context(mpl.rcParamsDefault):
        fig = plt.figure(figsize=(6, 1), frameon=False)
        ax = fig.add_subplot(111)
        ax.pcolor(np.linspace(1, ncolor, ncolor).reshape(1, ncolor), cmap=cm)
        xt = ax.set_xticks([])
        yt = ax.set_yticks([])
    return fig

For a list of colors

This is how to use the function :

colors = ["#9b59b6", "#3498db", "#95a5a6", "#e74c3c", "#34495e", "#2ecc71"]

And you will simply get a square for each color :


The source code
import matplotlib.pyplot as plt
def show_colors(colors):
    Draw a square for each color contained in the colors list
    given in argument.
    with plt.rc_context(plt.rcParamsDefault):
        fig = plt.figure(figsize=(6, 1), frameon=False)
        ax = fig.add_subplot(111)
        for x, color in enumerate(colors):
                    (x, 0), 1, 1, facecolor=color
        ax.set_xlim((0, len(colors)))
        ax.set_ylim((0, 1))
    return fig

Color scale and color maps with python

Treemaps plot using matplotlib and python

Here is an example of a treemap plot using python, matplotlib and the squarify module.

The data comes from the open data of the Communauté d’Agglomération Pau-Pyrénées.

The idea is the following :

  • The area of each square is the surface area of the towns
  • The color scale is over the population in 2011

Source code :

Gist code here !



Other version

Here is an other version where Pau is present. Nevertheless, I draw a white square background because the population of Pau is huge compared to other towns and thus I did not include it in the color scale.


Bands diagram using VASP and pymatgen

This article presents a python source code in order to plot the bands diagram of graphene calculated using VASP. The plot is done using pymatgen library and RGB coloring adapted from this example. Here is the output obtained with the script :

Graphene bands diagram
Graphene bands diagram

On the band diagram, the contribution of s, (px,py) and pz atomic orbital is map on a RGB scale. Red is associated to a s contribution, green to (px, py) contribution and blue to pz contribution. The same function or procedure can be used to map atomic contributions on the band diagram.

At the end of the article you can find a tarball containing all VASP input/ouput files and the python source code. This calculation is not accurate and is used only for band diagram plotting illustration. For example, you can see a short gap on the DOS which does not exist.

Edit 2016, March 15, update : source code is now python3 and pymatgen version 3.3.5 compatible. You can find the below scripts in an up to date version on github : bandstructureplots

#!/usr/bin/env python
# -*- coding=utf-8 -*-
import sys
import numpy as np
from numpy import array as npa
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.gridspec import GridSpec
import pymatgen as mg
from pymatgen.io.vasp.outputs import Vasprun, Procar
from pymatgen.symmetry.bandstructure import HighSymmKpath
from pymatgen.electronic_structure.core import Spin, Orbital
def rgbline(ax, k, e, red, green, blue, alpha=1.):
    # creation of segments based on
    # http://nbviewer.ipython.org/urls/raw.github.com/dpsanders/matplotlib-examples/master/colorline.ipynb
    pts = np.array([k, e]).T.reshape(-1, 1, 2)
    seg = np.concatenate([pts[:-1], pts[1:]], axis=1)
    nseg = len(k) - 1
    r = [0.5 * (red[i] + red[i + 1]) for i in range(nseg)]
    g = [0.5 * (green[i] + green[i + 1]) for i in range(nseg)]
    b = [0.5 * (blue[i] + blue[i + 1]) for i in range(nseg)]
    a = np.ones(nseg, np.float) * alpha
    lc = LineCollection(seg, colors=list(zip(r, g, b, a)), linewidth=2)
if __name__ == "__main__":
    # read data
    # ---------
    # kpoints labels
    path = HighSymmKpath(mg.Structure.from_file("./opt/CONTCAR")).kpath["path"]
    labels = [r"$%s$" % lab for lab in path[0][0:4]]
    # bands
    bands = Vasprun("./bands/vasprun.xml").get_band_structure("./bands/KPOINTS", line_mode=True)
    # projected bands
    data = Procar("./bands/PROCAR").data
    # density of state
    dosrun = Vasprun("./dos/vasprun.xml")
    # set up matplotlib plot
    # ----------------------
    # general options for plot
    font = {'family': 'serif', 'size': 24}
    plt.rc('font', **font)
    # set up 2 graph with aspec ration 2/1
    # plot 1: bands diagram
    # plot 2: Density of State
    gs = GridSpec(1, 2, width_ratios=[2, 1])
    fig = plt.figure(figsize=(11.69, 8.27))
    fig.suptitle("Bands diagram of graphene")
    ax1 = plt.subplot(gs[0])
    ax2 = plt.subplot(gs[1])  # , sharey=ax1)
    # set ylim for the plot
    # ---------------------
    emin = 1e100
    emax = -1e100
    for spin in bands.bands.keys():
        for b in range(bands.nb_bands):
            emin = min(emin, min(bands.bands[spin][b]))
            emax = max(emax, max(bands.bands[spin][b]))
    emin -= bands.efermi + 1
    emax -= bands.efermi - 1
    ax1.set_ylim(emin, emax)
    ax2.set_ylim(emin, emax)
    # Band Diagram
    # ------------
    # sum up contribution over carbon atoms
    data = data[Spin.up].sum(axis=2)
    # sum up px and py contributions and normalize contributions
    contrib = np.zeros((bands.nb_bands, len(bands.kpoints), 3))
    for b in range(bands.nb_bands):
        for k in range(len(bands.kpoints)):
            sc = data[k][b][Orbital.s.value]**2
            pxpyc = data[k][b][Orbital.px.value]**2 + \
            pzc = data[k][b][Orbital.pz.value]**2
            tot = sc + pxpyc + pzc
            if tot != 0.0:
                contrib[b, k, 0] = sc / tot
                contrib[b, k, 1] = pxpyc / tot
                contrib[b, k, 2] = pzc / tot
    # plot bands using rgb mapping
    for b in range(bands.nb_bands):
                [e - bands.efermi for e in bands.bands[Spin.up][b]],
                contrib[b, :, 0],
                contrib[b, :, 1],
                contrib[b, :, 2])
    # style
    ax1.set_ylabel(r"$E - E_f$   /   eV")
    # fermi level at 0
    ax1.hlines(y=0, xmin=0, xmax=len(bands.kpoints), color="k", lw=2)
    # labels
    nlabs = len(labels)
    step = len(bands.kpoints) / (nlabs - 1)
    for i, lab in enumerate(labels):
        ax1.vlines(i * step, emin, emax, "k")
    ax1.set_xticks([i * step for i in range(nlabs)])
    ax1.set_xlim(0, len(bands.kpoints))
    # Density of state
    # ----------------
    ax2.set_xticks(np.arange(0, 1.5, 0.4))
    ax2.set_xticklabels(np.arange(0, 1.5, 0.4))
    ax2.set_xlim(1e-6, 1.5)
    ax2.hlines(y=0, xmin=0, xmax=1.5, color="k", lw=2)
    ax2.set_xlabel("Density of State")
    # s contribution
    ax2.plot(npa(dosrun.pdos[0][Orbital.s][Spin.up]) +
             dosrun.tdos.energies - dosrun.efermi,
             "r-", label="s", linewidth=2)
    # px py contribution
    ax2.plot(npa(dosrun.pdos[0][Orbital.px][Spin.up]) +
             npa(dosrun.pdos[1][Orbital.px][Spin.up]) +
             npa(dosrun.pdos[0][Orbital.py][Spin.up]) +
             dosrun.tdos.energies - dosrun.efermi,
             label="(px, py)",
    # pz contribution
    ax2.plot(npa(dosrun.pdos[0][Orbital.pz][Spin.up]) +
             dosrun.tdos.energies - dosrun.efermi,
             "b-", label="pz", linewidth=2)
    # total dos
                     dosrun.tdos.energies - dosrun.efermi,
                     color=(0.7, 0.7, 0.7),
                     facecolor=(0.7, 0.7, 0.7))
             dosrun.tdos.energies - dosrun.efermi,
             color=(0.6, 0.6, 0.6),
             label="total DOS")
    # plot format style
    # -----------------
    ax2.legend(fancybox=True, shadow=True, prop={'size': 18})
    # plt.show()
    plt.savefig(sys.argv[0].strip(".py") + ".pdf", format="pdf")


Hereafter there are two more examples for Cu and Si. The script may be simpler as all data can be extracted using pymatgen methods.



Bands diagram of copper

Bands diagram of silicon

Outils pour la symétrie moléculaire : théorie des groupes

En chimie et en physique, on utilise les propriétés de symétrie du système pour simplifier la résolution. L’outils mathématiques à la base de l’utilisation de la symétrie en chimie est la théorie des groupes. Le livre « chimie et théorie des groupes » de Paul H. Walton (édition De Boeck) est un ouvrages bien adapté aux chimistes ou physico-chimistes qui désirent découvrir cette approche.

Lorsqu’on utilise la théorie des groupes, une étape primordiale et très mécanique est la décomposition de la représentation d’un groupe en représentation irréductibles de ce groupe (analogue à déterminer les coordonnées d’un vecteur dans une base orthonormée). Le module python ci-dessous permet d’obtenir automatiquement la décomposition. Voici un exemple :

>>> from groupes import c4v
>>> rep = [4, 0, 0, 2, 0]
>>> print(c4v)
 C4v  |     1E      2C4      1C2    2sigma_v 2sigma_d
  A1  |     1        1        1        1        1    
  A2  |     1        1        1        -1       -1   
  B1  |     1        -1       1        1        -1   
  B2  |     1        -1       1        -1       1    
  E   |     2        0        -2       0        0
>>> c4v.ordre
>>> c4v.reduce(rep)
1 A1 + 1 B1 + 1 E

Je n’ai pas entré tous les groupes ponctuels de symétrie mais uniquement ceux dont j’ai eu besoin pour le moment. Les groupes disponnibles sont :

  • C2v
  • C3v
  • C4v
  • D3h
  • D4h


vasptools : a package for vasp post treatments

Since may 2014, I did not continue the development of vasptools. I strongly recommend you tu use pymatgen instead.

vasptools is a package wrote in python in order to do pre or post treatments of a VASP calculations. This package contains :

  • python modules : vasptools, crystal, myxml, adsorption
  • several script using vasptools module

Scripts and functions read the data into the vasprun.xml file because it contains both input parameters and the results of a calculations. The main features are :

  • extract or plot density of state
  • extract or plot bands
  • get structural data
  • control convergence
  • a tools to put a molecule at a given position on a slab
  • some simple operation on CHGCAR file (split, sum)

You can find the whole documentation here : vasptools documentation.

vasptools.tar.gz (version 58 : 20/03/2013)

If you want to contribute, please contact me.

data2file : write block data into a file with python


This module provide the function d2f. The aim of d2f is to easily write block data into a
file with a given format and in one line. Maybe d2f looks like the fortran command write.

The up to date version is now on github : data2file

The main advantage of d2f is that you can give the format in a compact way and it avoids redondant
syntax. The syntax for format specification is the folowing :


as an example is clearer, this syntax


means that you want d2f to write two integers on four columns, then four strings on ten
columns and then height floating numbers on height columns with three digit after the comma.


 d2f(blockdata [, name[, fmt[, header[, overwrite]]]])
  • blockdata (list) : list of data, each element of the list in the first direction will be written on the same line.
  • name (string) : output file name. This could be a path
  • fmt (string) : format specification, example : ‘2x4d;4x10s;8×8.3f’
  • header (string) : will be writtend at the top of the output file
  • overwrite (bool) : if True, if the output file exists it will be overwritten

An example

Assume you have three list objects, l1, l2 and l3 and you want to print l1 on the first column, l2 on the second column and l3 on the third column of an output file called "out.dat". First you have to create a block data object. For that the easiest is to use zip(l1, l2, l3). Assume l1 is an integer, and l2 and l3 are floating numbers. You will use fmt="1x5d;2xf10.4" in order to write l1 as an integer on 5 columns and l2 and l3 as floats on 10 column with 4 digit after the comma. Thus simpliest syntax is :

>>> d2f(zip(l1, l2, l3), fmt = "1x5d;2xf10.4")

If you want to precise a file name :

>>> d2f(zip(l1, l2, l3), name = "output.dat", fmt = "1x5d;2xf10.4")

Python et (g)vim

Je donne ici une configuration et des plugins adaptés à l’écriture de code python avec vi. J’utilise la version graphique de vi : gvim. Ce que je propose ici n’est donc peut être pas commode en pratique avec la version terminal de vi.

plugins :

Voici les plugins vi que j’utilise :

  • comments.vim : permet de commenter ou décommenter une ligne ou un groupes de lignes avec les touches ctrl+C ctrl+X. Ce plugin fonctionne avec tout une liste de langage.
  • jpythonfold.vim : ajoute automatiquement des replis pour les classes / méthodes / fonctions de pythons.
  • pyflakes.vim : vérification au fil de l’eau de la syntaxe python. pyflakes agit à la manière d’un correcteur orthographique et souligne en rouge les lignes de codes incorrectes, les variables inutilisées ou non assignées …

Installation :

Dans un premier temps on installe comments.vim et jpythonfold.vim. S’ils n’existent pas créer les dossiers .vim dans votre dossier utilisateur et un sous dossier plugin

mkdir -pv $HOME/.vim/plugin

Copier simplement les fichiers comments.vim et jpythonfold.vim dans le dossier .vim/plugin.

Pour jpythonfold, il est pratique de ne l’activer que lorsque le fichier est un fichier python. Pour cela l’auteur propose de décommenter certaines lignes dans le fichier jpythonfold.vim. C’est aux alentours de la ligne 40. Chez moi ça n’a pas fonctionné. En m’inspirant de ce qui est fait dans comments.vim, je vous propose de mettre ces lignes à la suite de celles proposées en laissant celles qui étaient proposées par l’auteur commentées.

let file_name = buffer_name("%")
if file_name =~ '\.py$'

C’est assez simple à comprendre, si le nom de votre fichier se termine par .py alors le plugin est activé.

Pour installer pyflakes, la procédure est expliquée sur la page du plugin. S’ils n’existent pas créer les dossiers ftplugin et python respectivement dans les dossiers $HOME/.vim et ftplugin :

mkdir -pv $HOME/.vim/ftplugin/python

Après avoir téléchargé l’archive de pyflakes, extraire le contenu dans $HOME/.vim/ftplugin/python. Il faut ensuite activer les ftplugin (pour File Type Plugin) qui assure le chargement du plugins en fonction du fichier considéré. S’il n’existe pas, créer le fichier .vimrc dans votre dossier utilisateur et ajouter la ligne suivante à l’intérieur :

filetype plugin indent on

Configuration supplémentaires de vi

Le fichier .vimrc permet de configurer (g)vim à votre gout. Voici quelques options spécifiques à la programmation qui peuvent être intéressantes :

syntax on               " coloration syntaxique
set showmatch        " affiche les paires de parenthese
" pour l'indentation
set autoindent		" copy indent from the current line to the new line
set smartindent		" automatically indent a line depending on previous line
set expandtab           " remplace les tab par des espaces
set shiftwidth=4        " indentation = 4 espaces
" encodage
set encoding=utf-8
" une abreviation
ab pcom # * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *

Pour plus d’options de configurations de (g)vim, vous pouvez voir cet article.

Bonne programmation !

Guide pour tracer des graphiques en coordonnées polaires ou sphériques

Pour tracer un graphique en utilisant les coordonnées polaires le quadrillage habituel (feuille à carreaux, papier millimétré …) n’est pas adapté car il est fait pour des coordonnées cartésiennes. Les codes des figures que je donne ici permettent de tracer des guides qui facilitent le tracé de courbes en coordonnées polaire.

Guide pour tracer des courbes en coordonnées polaires

Lorsqu’on utilise les coordonnées sphériques, on peut choisir de tracer la fonction dans un plan de l’espace, par exemple le plan (xOz). En coordonnées sphériques, l’angle \theta n’est pas défini de la même manière qu’en coordonnées polaires. Il faut donc faire quelques modifications. Voici ci-dessous un exemple de tracé de la partie angulaire de l’orbitale atomique d_{z^2} dans le plan (xOz) en coordonnées sphériques.

orbitale atomique dz2 tracée avec tikz

Le fichier polar_graph.tex ci-dessous contient :

  • Le guide pour tracer des courbes en coordonnées polaires
  • Le guide pour tracer des courbes en coordonnées sphériques dans le plan (xOz)
  • L’exemple de dessin de la partie angulaire de l’orbitale atomique d_{z^2}

Fichier source : polar_graph.tex

Fichier pdf : polar_graph.pdf

Ci-dessous un script en python qui permet de faire un graphique en coordonnées polaires en utilisant la librairies matplotlib. Le défaut est que c’est un graphique en coordonnées polaires, donc du fait de la définition de l’angle \theta, l’axe z est horizontal.

orbitale atomique dz2 tracée avec python

script : dz2.py

bigfiles.py : rechercher les plus gros fichiers dans une arborescence de dossier

Description :

Voici un script écrit en python qui recherche les plus gros fichiers dans une arborescence de dossiers de façon récursive. Par défaut il affiche les vingt plus gros fichiers. On peut indiquer au script le chemin d’un dossier à analyser ou encore le nombre de fichiers que l’on veut afficher. Si on donne un nombre de fichiers négatifs le script affichera la liste de tous les fichiers (attention elle peut être longue). Quelques exemples sont donnés dans la doc, vous pouvez passer l’option -h ou –help à l’exécution pour l’afficher. Tous les fichiers sont pris en considérations, y compris les fichiers cachés (commençant par un . sous linux).

Installation (sous linux) :

Téléchargez le fichier ci-dessous et donnez lui le nom que vous souhaitez, par exemple bigfiles. Placez ensuite ce fichier dans un dossier contenu dans votre PATH comme par exemple /usr/local/bin si vous avez les droits administrateur. Il ne faut pas oublier de rendre le fichier exécutable :
chmod +x bigfile. Vous pourrez ensuite exécuter bigfile depuis une console linux.

script :


# -*- coding=utf-8 -*-
print the first bigest files in all subdirectories of a given path
        bigfiles [OPTIONS] ... [PATH]
        Print the first 20 bigest files in directory [PATH] and in all subdirectories of
        [PATH]. Default [PATH] is the current working directory.
        -h, --help
            print this help
        -n, --nfiles=N
            output the first N bigest files, instead if the first 20. If N is negative,
            print all files.
        bigfile -h
        bigfile -n 50
        bigfile ./another/directory/
        bigfile -n -1 
        bigfile -n 20 ./anotherdirectory
__licence__ = "GPL"
__author__ = "Germain Vallverdu <germain .vallverdu@univ-pau.fr>"
import doctest
import os
import sys
MEGA_OCTET = 1048576
GIGA_OCTET = 1073741824
TERA_OCTET = 1099511627776
def humanSize(taille):
    """ return a byte number in human readable format (1ko 234Mo 2Go
    >>> humanSize( 345)
    '345 o'
    >>> humanSize( 123478)
    '120.6 ko'
    >>> humanSize( 983435743)
    '937.9 Mo'
    >>> humanSize( 12983435743)
    '12.1 Go'
    >>> humanSize( 755812983435743)
    '687.4 To'
    taille = float(taille)
    if taille >= TERA_OCTET:
        size = "%.1f To" % (taille / TERA_OCTET)
    elif taille >= GIGA_OCTET:
        size = "%.1f Go" % (taille / GIGA_OCTET)
    elif taille >= MEGA_OCTET:
        size = "%.1f Mo" % (taille / MEGA_OCTET)
    elif taille >= KILO_OCTET:
        size = "%.1f ko" % (taille / KILO_OCTET)
        size = "%.0f o" % taille
    return size
def BigFiles( dossier = "./", nmax = 20):
    """ print the first bigest files in all subdirectories of a given path 
        arguments :
            * dossier : path where you want to list bigest files
            * nmax    : maximum number of files you want BigFiles to list
    fichier = list()
    tailleTotale = 0
    nFiles = 0
    for root, dirs, files in os.walk(dossier):
        for name in files:
            if not os.path.isfile(os.path.join(root,name)): continue
            taille = os.path.getsize(os.path.join(root, name))
            nFiles += 1
            tailleTotale += taille
            fichier.append({"nom":name, "dossier":root, "taille":taille})
    print("\n" + " * * * BigFiles * * *".center(80) + "\n")
    print(" Analyse du dossier         : " + dossier)
    print(" Nombre de fichiers traités : " + str(nFiles))
    print(" Taille totale des fichiers : " + humanSize(tailleTotale))
    # sort files and print the first nmax
    sorted_fichier = sorted( fichier, key = lambda f:f["taille"], reverse = True)
    for i, f in enumerate(sorted_fichier):
        print(str(i+1).rjust(6) + humanSize(f["taille"]).rjust(10) + f["nom"].rjust(30) + "   " + f["dossier"])
        if i == (nmax-1): break
if __name__ == "__main__":
    # --------------------------------------------------------
    # default values
    # --------------------------------------------------------
    dossier = "./"
    nmax = 20
    # --------------------------------------------------------
    # get options
    # --------------------------------------------------------
    narg = len(sys.argv)
    if narg > 1:
        args = sys.argv
        i = 1
        while i < narg:
            if args[i] == "-h" or args[i] == "--help":
            elif args[i] == "-n" or args[i] == "--nfiles":
                    nmax = int(args[i+1])
                except ValueError:
                    print("\nError : bad arguments " + args[i])
                    print("    try : " + args[0] + " --help\n")
                except IndexError:
                    print("\nError : bad arguments " + args[i])
                    print("    try : " + args[0] + " --help\n")
                i += 2
            elif args[i] == "-v":
                i += 1
                if i == narg - 1:
                    dossier = args[i]
                    print("\nError : bad arguments")
                    print("    try : " + args[0] + " --help\n")
                i += 1
    BigFiles(dossier, nmax)

Installation de modules python supplémentaires

Supposons que vous avez écris un module python contenant des classes ou des fonctions ou que vous en avez téléchargé un sur internet. Où devez vous le placer pour y avoir accès lorsque vous exécutez un programme python ou que vous utilisez le prompt de python ?

Lorsqu’on importe un module python avec l’instruction import ou from, python cherche le module correspondant dans une liste d’emplacements (de dossiers) défini dans la variable d’environnement PYTHONPATH. On peut connaître la liste de ces emplacements grâce à la variable path du module sys :

[Ger@kiwi:/home/Ger]> python
Python 2.6.6 (r266:84292, Sep 15 2010, 16:22:56) 
[GCC 4.4.5] on linux2
Type "help", "copyright", "credits" or "license" for more information.
>>> import sys
>>> for p in sys.path:
...     print(p)

La variable sys.path est en fait une liste python. On voit, entre autre, que le dossier de travail (ici /home/Ger/) est contenu dans la variable sys.path. Une première solution pour importer dans un code python un module supplémentaire est donc de le placer dans le dossier de travail. Cependant cette solution n’est pas très agréable étant donné qu’il faudra recopier le module dans tous les dossiers ou vous en aurez besoin. Une seconde solution est de rajouter des éléments à la liste sys.path avant d’importer les modules :

import sys
# importation de votre module
import mon_module

De façon plus automatique on peut ajouter des emplacements à la variable d’environnement PYTHONPATH. Voici quelques lignes en bash que vous pouvez ajouter à votre fichier .bashrc pour ajouter un ensemble de chemins de façon automatique. Ce petit script suppose que tous vos modules python supplémentaires sont enregistrés dans un même dossier et dans les sous dossier de ce dernier. Ce dossier nommé classDir est défini sur la première ligne du script qui ajoute simplement à la variable PYTHONPATH le dossier classDir et l’ensemble des dossiers qu’il contient. Voici le script :

# python : liste des dossiers contenant vos classes python
liste=`ls $classDir`
for dossier in $liste