Phase plot using matplotlib tricontourf

Question

I want to plot an image of the results of a finite element simulation with a personalized colormap. I have been trying to use tricontourf to plot it as follow :

    #Z = self.phi.compute_vertex_values(self.mesh)
    Z = np.mod(self.phi.compute_vertex_values(self.mesh),2*np.pi)
    triang = tri.Triangulation(*self.mesh.coordinates().reshape((-1, 2)).T,
                               triangles=self.mesh.cells())
    zMax = np.max(Z)
    print(zMax)

    #Colormap creation
    nColors = np.max(Z)*200/(2*np.pi)
    phiRange = np.linspace(0,zMax,nColors)
    intensity = np.sin(phiRange)**2
    intensityArray = np.array([intensity, intensity, intensity])
    colors = tuple(map(tuple, intensityArray.T))
    self.cm = LinearSegmentedColormap.from_list("BAM", colors, N=nColors)

    #Figure creation
    fig, ax = plt.subplots()    
    levels2 = np.linspace(0., zMax,nColors)

    cax = ax.tricontourf(triang, Z,levels=levels2, cmap = self.cm) #plot of the solution
    fig.colorbar(cax)

    ax.triplot(triang, lw=0.5, color='yellow') #plot of the mesh

    plt.savefig("yolo.png")
    plt.close(fig)

And it gives the result :

As you can see there are some trouble where the phase goes from 2pi to 0 that comes from tricontourf when there is a modulo...

My first idea for work around was to work directly on my phase Z. The problem is that if I do this I need to create a much larger colormap. Ultimately, the phase will be very large and so will be the colormap if I want a correct color resolution... Furthemore I would like to have only one period in the colormap on the right (just like in the first figure).

Any idea how I could obtain a figure just like the second one, with a colormap just like the one from the first figure and without creating a very large and expensive colormap ?

EDIT : I have written a small code that is runnable out of the box : It reproduces the problem I have and I have also tried to apply Thomas Kuhn answer to my preoblem. However, it seems that there are some problem with the colorbar... Any idea how I could fix this ?

import matplotlib.pyplot as plt
import matplotlib.tri as mtri
import numpy as np
import matplotlib.colors as colors

class PeriodicNormalize(colors.Normalize):
    def __init__(self, vmin=None, vmax=None, clip=False):
        colors.Normalize.__init__(self, vmin, vmax, clip)    

    def __call__(self, value, clip=None):
        x, y = [self.vmin, self.vmax], [0, 1]    

        return np.ma.masked_array(np.interp(
            np.mod(value-self.vmin, self.vmax-self.vmin),x,y
        ))

# Create triangulation.
x = np.asarray([0, 1, 2, 3, 0.5, 1.5, 2.5, 1, 2, 1.5])
y = np.asarray([0, 0, 0, 0, 1.0, 1.0, 1.0, 2, 2, 3.0])
triangles = [[0, 1, 4], [1, 2, 5], [2, 3, 6], [1, 5, 4], [2, 6, 5], [4, 5, 7],
             [5, 6, 8], [5, 8, 7], [7, 8, 9]]
triang = mtri.Triangulation(x, y, triangles)
cm = colors.LinearSegmentedColormap.from_list('test', ['k','w','k'], N=1000) 

#Figure 1 : modulo is applied on the data : 
#Results : problem with the interpolation, but the colorbar is fine
z = np.mod(10*x,2*np.pi)
zMax = np.max(z)
levels = np.linspace(0., zMax,100)

fig1, ax1 = plt.subplots()
cax1=ax1.tricontourf(triang, z,cmap = cm,levels= levels)
fig1.colorbar(cax1)
plt.show()

#Figure 2 : We use the norm parameter with a custom norm that does the modulo
#Results : the graph is the way it should be but the colormap is messed up
z = 10*x
zMax = np.max(z)
levels = np.linspace(0., zMax,100)

fig2, ax2 = plt.subplots()
cax2=ax2.tricontourf(triang, z,levels= levels,norm = PeriodicNormalize(0, 2*np.pi),cmap = cm)
fig2.colorbar(cax2)
plt.show()

Last solution would be to do as I did above : to create a much larger colormap that goes up to zmax and is periodic every 2 pi. However the colorbar would not be nice...

here are the results :

Answer 1

I'm guessing that your problem arises from using modulo on your data before you call tricontourf (which, I guess, does some interpolation on your data and then maps that interpolated data to a colormap). Instead, you can pass a norm to your tricontourf function. Writing a small class following this tutorial, you can make the norm take care of the modulo of your data. As your code is not runnable as such, I came up with an a bit simpler example. Hopefully this is applicable to your problem:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors   

class PeriodicNormalize(colors.Normalize):
    def __init__(self, vmin=None, vmax=None, clip=False):
        colors.Normalize.__init__(self, vmin, vmax, clip)    

    def __call__(self, value, clip=None):
        x, y = [self.vmin, self.vmax], [0, 1]    

        return np.ma.masked_array(np.interp(
            np.mod(value-self.vmin, self.vmax-self.vmin),x,y
        ))

fig,ax = plt.subplots()
x,y = np.meshgrid(
    np.linspace(0, 1, 1000),
    np.linspace(0, 1, 1000),
)

z = x*10*np.pi    
cm = colors.LinearSegmentedColormap.from_list('test', ['k','w','k'], N=1000)    
ax.pcolormesh(x,y,z,norm = PeriodicNormalize(0, 2*np.pi), cmap = cm)
plt.show()

The result looks like this:

EDIT:

As the ContourSet you get back from tricontourf spans the full phase, not just the first [0,2pi], the colorbar is created for that full range, which is why you see the colormap repeat itself many times. I'm not quite sure if I understand how the ticks are created, but I'm guessing that it would be quite some work to get that automated to work right. Instead, I suggest to generate a colorbar "by hand", as is done in this tutorial. This, however, requires that you create the axes (cax) where the colorbar is put yourself. Luckily there is a function called matplotlib.colorbar.make_axes() that does this for you (all thanks goes to this answer). So, instead of your original colorbar command, use these two lines:

cax,kw = mcbar.make_axes([ax2], location = 'right')
cb1 = mcbar.ColorbarBase(cax, cmap = cm, norm = norm, orientation='vertical')

To get this picture: