matplotlib shift pcolormesh plot to symmetrized coordinates

Question

I have some 2D data with x and y coordinates both within [0,1], plotted using pcolormesh. Now I want to symmetrize the plot to [-0.5, 0.5] for both x and y coordinates. In Matlab I was able to achieve this by changing x and y from e.g. [0, 0.2, 0.4, 0.6, 0.8] to [0, 0.2, 0.4, -0.4, -0.2], without rearranging the data. However, with pcolormesh I cannot get the desired result.

A minimum example is shown below, with data represented simply by x+y:

import matplotlib.pyplot as plt
import numpy as np

x,y = np.mgrid[0:1:5j,0:1:5j] 

fig,(ax1,ax2,ax3) = plt.subplots(1,3,figsize=(9,3.3),constrained_layout=1)
# original plot spanning [0,1]
img1 = ax1.pcolormesh(x,y,x+y,shading='auto')

# shift x and y from [0,1] to [-0.5,0.5]
x = x*(x<0.5)+(x-1)*(x>0.5)
y = y*(y<0.5)+(y-1)*(y>0.5)
img2 = ax2.pcolormesh(x,y,x+y,shading='auto') # similar code works in Matlab

# for this specific case, the following is close to the desired result, I can just rename x and y tick labels
# to [-0.5,0.5], but in general data is not simply x+y
img3 = ax3.pcolormesh(x+y,shading='auto') 
fig.colorbar(img1,ax=[ax1,ax2,ax3],orientation='horizontal')

The corresponding figure is below, any suggestion on what is missed would be appreciated!

@JodyKlymak I just double checked that in Matlab I was doing a 3D version of the same problem and used isosurface, which worked. Using pcolor in Matlab indeed is not enough. Naively I was thinking to effectively rearrange the "Z" data by moving its X and Y coordinates, but maybe that only works for scatter plots... — Phyinmi

scleronomic scleronomic · Accepted Answer · 2020-09-21T07:23:33

Let's look at what you want to achieve in a 1D example. You have x values between 0 and 1 and a dummy function f(x) = 20*x to produce some values.

# x  = [0, .2, .4, .6, .8] -> [0, .2, .4, -.4, -.2] -> [-.4, .2, .0, .2, .4])
# fx = [0,  4,  8, 12, 16] -> [0,  4,  8,  12,  16] -> [ 12, 16,  0,  4,  8]
#                                          ^ only flip and shift x not fx ^

You could use np.roll() to achieve the last operation. I used n=14 to make the result better visible and show that this approach works for arbitrary n.

import numpy as np
import matplotlib.pyplot as plt

n = 14
x, y = np.meshgrid(np.linspace(0, 1, n, endpoint=False),
                   np.linspace(0, 1, n, endpoint=False))

z = x + y

x_sym = x*(x <= .5)+(x-1)*(x > .5)
# array([[ 0. ,  0.2,  0.4, -0.4, -0.2], ...
x_sym = np.roll(x_sym, n//2, axis=(0, 1))
# array([[-0.4, -0.2,  0. ,  0.2,  0.4], ...

y_sym = y*(y <= .5)+(y-1)*(y > .5)
y_sym = np.roll(y_sym, n//2, axis=(0, 1))

z_sym = np.roll(z, n//2, axis=(0, 1))
# array([[1.2, 1.4, 0.6, 0.8, 1. ],
#        [1.4, 1.6, 0.8, 1. , 1.2],
#        [0.6, 0.8, 0. , 0.2, 0.4],
#        [0.8, 1. , 0.2, 0.4, 0.6],
#        [1. , 1.2, 0.4, 0.6, 0.8]])

fig, (ax1, ax2) = plt.subplots(1, 2)
img1 = ax1.imshow(z, origin='lower', extent=(.0, 1., .0, 1.))
img2 = ax2.imshow(z_sym,  origin='lower', extent=(-.5, .5, -.5, .5))

matplotlib shift pcolormesh plot to symmetrized coordinates

1 Answers