Plot Imshow on a Torus (lattice with PBC)

Is there any "simple" solution in matplotlib or another Python package to plot a square lattice on a 2d torus (aka lattice with periodic boundary conditions)?

Assume i have a simple 2D array

# ...
a = np.random.random((50, 50))
plt.imshow(a)

I would like to wrap this plane into a torus, which can be achieved e.g. with

from mpl_toolkits.mplot3d import Axes3D
# Generating Torus Mesh
angle = np.linspace(0, 2 * np.pi, 100)
theta, phi = np.meshgrid(angle, angle)
r, R = .25, 1.
X = (R + r * np.cos(phi)) * np.cos(theta)
Y = (R + r * np.cos(phi)) * np.sin(theta)
Z = r * np.sin(phi)


fig = plt.figure()
ax = fig.add_subplot(projection = '3d')

ax.set_xlim3d(-1, 1)
ax.set_ylim3d(-1, 1)
ax.set_zlim3d(-1, 1)
ax.plot_surface(X, Y, Z, rstride = 1, cstride = 1)

plt.show()

I thought about encoding somehow the information on X and Y value in a colormap to pass to plot_surface's option cmap. The colormap shall have each color of the image according to content of array a.

Ideas?

Your question will be generalized as 'is there any simple solution to map a 2D plane onto a 3D surface?' The Matplotlib third-party package, S3Dlib, may provide a relatively simple solution. This package allows direct functional or image mapping onto predefined mesh 3D surfaces. In cylindrical coordinates ( r,t,z ), functional mapping of a cylinder to a torus is:

def torusFunc(rtz) :
    r,t,z = rtz
    Z = np.sin(z*np.pi)/3
    R = r + np.cos(z*np.pi)/3
    return R,t,Z

Any face coordinate of N number of faces on a cylinder can be assigned a random value using:

def random_colors(rtz):
    r,t,z = rtz
    N = len(t)*len(z)
    return np.random.random(N)

Then using these two functions, a cylinder object is created, color assigned to each face, and then mapped to the toroidal geometry. The object is added to the Axes3D with the result:

For the above figure, shading and hignlighting were used to clarify the surface geometry. The code is:

import numpy as np
from matplotlib import pyplot as plt
import s3dlib.surface as s3d

# 1. Define functions .....................................
def torusFunc(rtz) :
    r,t,z = rtz
    Z = np.sin(z*np.pi)/3
    R = r + np.cos(z*np.pi)/3
    return R,t,Z

np.random.seed(0)
def random_colors(rtz):
    r,t,z = rtz
    N = len(t)*len(z)
    return np.random.random(N)

# 2. Create and map the surface ...........................
torus = s3d.CylindricalSurface.grid(30,100,'r')
torus.map_cmap_from_op( random_colors, "viridis" )
torus.map_geom_from_op(  torusFunc ) 

# 3. Construct axes, add surface object, & show ...........
minmax,ticks = (-1.33,1.33), (-1,-.5,0,.5,1)
fig = plt.figure(figsize=plt.figaspect(1))
ax = plt.axes(projection='3d')
ax.set(xlim=minmax,  ylim=minmax,  zlim=minmax,
       xticks=ticks, yticks=ticks, zticks=ticks )
ax.set_title(torus.name)

ax.add_collection3d(torus.shade().hilite(.5))

fig.tight_layout()
plt.show()

Image mapping can also be used similar to functional mapping. Since mapping is visually obscured using a random colored image, a earth surface image will be used to demonstrate this:

The image originated from the NASA visible earth catalog. The following figure shows image mapping using planar, polar, cylindrical and spherical coordinate systems.

where the code is:

import numpy as np
from matplotlib import pyplot as plt
import s3dlib.surface as s3d

# 1. Define functions to examine ....................................
def planarFunc(xyz) : # Himmelblau
    x,y,z = xyz
    X,Y = 5*x, 5*y
    Z1 = np.square( X*X + Y - 11 )
    Z2 = np.square( Y*Y + X - 7  ) 
    Z = (Z1 + Z2)/500 - 1
    return x,y,Z

def screwFunc(rtz) :
    r,t,z = rtz
    T = 2*t
    Z = 0.2*(T - 2*np.pi)
    return r,T,Z

def torusFunc(rtz) :
    r,t,z = rtz
    Z = np.sin(z*np.pi)/3
    R = r + np.cos(z*np.pi)/3
    return R,t,Z

def swirlFunc(rtp) :
    r,t,p = rtp
    R = 0.4*(2 + np.sin(5*t + 3*p))
    return R,t,p
 
# 2. Setup and map surfaces .........................................
img_fname, rez = 'data/earth.png', 5

planar = s3d.PlanarSurface(rez,basetype='oct1')
planar.map_color_from_image(img_fname)
planar.map_geom_from_op(planarFunc)

screw = s3d.PolarSurface(rez, basetype='hex_s')
screw.map_color_from_image(img_fname)
screw.map_geom_from_op( screwFunc)

torus = s3d.CylindricalSurface(rez, basetype='squ_s')
torus.map_color_from_image(img_fname)
torus.map_geom_from_op(  torusFunc) 

swirl = s3d.SphericalSurface(rez)
swirl.map_color_from_image(img_fname)
swirl.map_geom_from_op( swirlFunc )

# 3. Construct figure, add surfaces, and plot ......................
minmax = (-.9,.9)
fig = plt.figure(figsize=plt.figaspect(1) )
for i,surface in enumerate( [planar,screw,torus,swirl] ) :
    ax = fig.add_subplot(221+i, projection='3d')
    ax.set(xlim=minmax, ylim=minmax, zlim=minmax)
    ax.set_title(surface.name)
    ax.set_axis_off()

    ax.add_collection3d(surface.shade().hilite(.5))

fig.tight_layout()
plt.show()

This technique may also be used with parametric functions. For example:

where the code is:

import numpy as np
from matplotlib import pyplot as plt
import s3dlib.surface as s3d

# 1. Define function to examine ....................................
def klein_bottle(rtz) :
    r,t,z = rtz
    u = np.pi*(1-z)/2
    v = -t
    cU, sU = np.cos(u), np.sin(u)
    cV, sV = np.cos(v), np.sin(v)
    x = -(2/15)*cU* \
        (  ( 3 )*cV + \
           ( -30 + 90*np.power(cU,4) - 60*np.power(cU,6) + 5*cU*cV )*sU  )
    y = -(1/15)*sU* \
        (  ( 3 - 3*np.power(cU,2) -48*np.power(cU,4) +48*np.power(cU,6) )*cV + \
           (-60 + ( 5*cU - 5*np.power(cU,3) - 80*np.power(cU,5) + 80*np.power(cU,7) )*cV  )*sU  )
    z = (2/15)*( 3 + 5*cU*sU )*sV
    return x,y,z

# 2. Setup and map surface .........................................
img_fname, rez = 'data/earth.png', 7

surface = s3d.CylindricalSurface(rez, color='w' )
surface.map_color_from_image(img_fname)
surface.map_geom_from_op( klein_bottle, returnxyz=True )
surface.transform(s3d.eulerRot(90,-90),translate=[0,0,2])

# 3. Construct figure, add surface plot ............................
minmax = (-1.5,1.5)
fig = plt.figure(figsize=plt.figaspect(1))
ax = plt.axes(projection='3d')
ax.set(xlim=minmax, ylim=minmax, zlim=minmax)
ax.set_title(surface.name)
ax.set_axis_off()
ax.view_init(elev=20, azim=-125)

ax.add_collection3d(surface.shade(.5,ax=ax).hilite(ax=ax))

fig.tight_layout()
plt.show()

The mathematical description of the geometry used in the code is from the Klein bottle Wikipedia page.