How to transform the coordinate axes in matplotlib

Question

I want to plot a hexagonal lattice starting off from the meshgrid function that gives me a square lattice.

import numpy as np
import matplotlib.pyplot as plt

xx, yy = np.meshgrid(np.arange(10), np.arange(10), indexing='ij')
plt.scatter(xx,yy)
plt.show()

If transform all of the points by the matrix A that contains the basis vectors of the hexagonal lattice I get:

A = np.array([[3./2, 3./2],[np.sqrt(3)/2, -np.sqrt(3)/2]])
pts = np.einsum('ij,jk->ik',A,np.array([xx.flatten(),yy.flatten()]))
plt.scatter(pts[0,:], pts[1,:])
plt.show()

I think it would be more elegant if, instead of transforming each point, I could just transform the axes of my coordinate system to get the same result. Is there a function in matplotlib that allows me to change the basis vectors from [1,0] and [0,1] to [3./2, sqrt(3)/2] and [3./2, -sqrt(3)/2]? Or can I give such a matrix to np.meshgrid?

Answer 1

You can provide an affine transform to the scatter to skew rotate the points.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.transforms as mtrans

xx, yy = np.meshgrid(np.arange(10), np.arange(10))
A = np.array([[3./2, 3./2, 0],[np.sqrt(3)/2, -np.sqrt(3)/2, 0], [0,0,1]])

fig, ax = plt.subplots()
plt.scatter(xx.flat, yy.flat, transform=mtrans.Affine2D(A) + ax.transData)
plt.show()