user10853
user10853

Reputation: 302

Python - Combine plots in grid

I wrote a function (see end) that calculates spherical harmonic coefficients for a specific order and degree and plots them on a sphere.

I would like to combine several of these spheres in a grid. I would like an end result similar to thisenter image description here

I tried with 2 plots using plt.subplots and plt.gridSpec to no avail. It always ends up putting the other plot outside. Here's the code I tried:

fig, axes = plt.subplots(ncols=1, nrows=2)
ax1, ax2 = axes.ravel()
ax1.plot(sh(6,6))
ax2.plot(sh(7,7))
plt.show()

I get the following figure:

enter image description here

and a traceback @ 3rd line:

ValueError: x and y must not be None

Also,

### GridSpec ####
plt.subplot2grid((2,2), (0,0))
sh(7,7)
plt.subplot2grid((2,2), (1, 0))
sh(8,7)
plt.subplot2grid((2,2), (1, 1))
sh(9,7)
plt.show()

results in 3 separate (not grid) plots.

It is a better result but the 3rd sphere should be on the right of the 2nd sphere unless I have done something wrong. Note: sh() is the function I wrote which calculates the spherical harmonics and plots the sphere with the spherical harmonics projections. In other words I have 2 spheres here. All I want to do is combine the two (actually more) spheres in a grid like the one above.

PS: I tried to work with Mayavi but I couldn't make it work. All the code on the website doesn't work for me. I will recheck it later but I am tight on time now so I wrote my own function.

The function I wrote:

def sh(l,m,cent_lat,cent_lon):
    # function that calculates the spherical harmonics of order l and degree m and visualizes it on a
    # sphere centered at (cent_lat, cent_lon) given in degrees

    if l < m:
        print "Order cannot be smaller than the degree! Try again."
    else:

        import numpy as np
        import scipy.special as sp
        from math import pi
        import matplotlib.cm as cm
        from mpl_toolkits.basemap import Basemap
        import matplotlib.pyplot as plt

        res = pi/100 # resolution
        theta = np.r_[0:2*pi:res]; phi = np.r_[0:pi:res] # theta: lon, phi: coalt

        coef = []
        for i in theta:
            for j in phi:
                coef.append(sp.sph_harm(m,l,i,j))
        coef = np.asarray(coef) # convert list to array
        coef = np.reshape(coef, (len(theta),-1)) # reshapte array as per number of angles


        ## Plotting ##

        # create lat/lon arrays
        lon = np.linspace(0,2*pi,len(theta))
        lat = np.linspace(-pi/2,pi/2,len(phi))
        colat = lat+pi/2 # colatitude array

        # create 2D meshgrid
        mesh_grid = np.meshgrid(lon, lat) # create a meshgrid out of lat/lon
        lon_grid = mesh_grid[0] # grab the meshgrid part for lon
        lat_grid = mesh_grid[1] # grab the meshgrid part for lat

        real_coef = np.real(coef) # read parts of the coefficients
        norm_coef = np.round(real_coef / np.max(real_coef),2) # normalize

        # set up orthographic map projection
        mp = Basemap(projection='ortho', lat_0 = cent_lat, lon_0 = cent_lon) # setup an orthographic basemap centered at lat_0 & lon_0
        # draw the edge of the map projection region (the projection limb)
        mp.drawmapboundary()

        # convert angles from radians to degrees & pipe them to basemap
        x,y = mp(np.degrees(lon_grid), np.degrees(lat_grid)) 

        cmap = cm.get_cmap('jet') # Set color map
        mp.pcolor(x,y,np.transpose(norm_coef), cmap=cmap)
        # cax = figure.add_axes([0.15,0.03,0.7,0.03])
#         cb = plt.colorbar(orientation = 'horizontal')
        plt.show()

Any help is appreciated.

Upvotes: 1

Views: 2926

Answers (1)

fjarri
fjarri

Reputation: 9726

GridSpec works for me (matplotlib v1.5.0 in case it matters):

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import LightSource
import matplotlib.cm as cm

import matplotlib.gridspec as gridspec

import numpy as np


u = np.linspace(0, 2 * np.pi, 100)
v = np.linspace(0, np.pi, 100)

x = 10 * np.outer(np.cos(u), np.sin(v))
y = 10 * np.outer(np.sin(u), np.sin(v))
z = 10 * np.outer(np.ones(np.size(u)), np.cos(v))

ls = LightSource(270, 45)
rgb = ls.shade(z, cmap=cm.gist_earth, vert_exag=0.1, blend_mode='soft')


def plot_sphere(s):

    s.plot_surface(x, y, z, rstride=4, cstride=4, facecolors=rgb,
                           linewidth=0, antialiased=False, shade=False)

    s.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
    s.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
    s.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))

    s.w_xaxis.line.set_color((1.0, 1.0, 1.0, 0.0))
    s.w_yaxis.line.set_color((1.0, 1.0, 1.0, 0.0))
    s.w_zaxis.line.set_color((1.0, 1.0, 1.0, 0.0))

    s.set_xticks([])
    s.set_yticks([])
    s.set_zticks([])


fig = plt.figure()

gs = gridspec.GridSpec(6, 7)

for i in range(6):
    for j in range(6):
        if i >= j:
            s = fig.add_subplot(gs[i,j+1], projection='3d')
            plot_sphere(s)


for i in range(6):

    pos = gs[i,0].get_position(fig)
    fig.text(pos.x0, (pos.y0 + pos.y1) / 2, "$\\ell = " + str(i) + "$", fontsize=16)

fig.savefig('t.png')

enter image description here

Upvotes: 3

Related Questions