Putting it together

Putting it together

import pandas as pd
import numpy as np

import matplotlib.pyplot as plt

%matplotlib notebook

px = pd.read_csv("px_2D.csv", header=None)

x = px.iloc[0, 1:]
y = px.iloc[1:, 0]
px_values = px.iloc[1:, 1:]

py = pd.read_csv("py_2D.csv", header=None)
py_values = py.iloc[1:, 1:]

pz = pd.read_csv("pz_2D.csv", header=None)
pz_values = pz.iloc[1:, 1:]
# Calculate min, max, levels, ticks
pmin = px_values.min().min()
pmax = px_values.max().max()
levels = np.linspace(pmin, pmax, 30)
ticks = np.linspace(pmin, pmax, 6)


# Create plots
fig, ax = plt.subplots(1,3, figsize=(9,4)) # Add figsize argument here to set figure side

ax[0].contourf(x, y, px_values, vmin=pmin, vmax=pmax, levels=levels, cmap="RdBu")
ax[1].contourf(x, y, py_values, vmin=pmin, vmax=pmax, levels=levels, cmap="RdBu")
CS= ax[2].contourf(x, y, pz_values, vmin=pmin, vmax=pmax, levels=levels, cmap="RdBu")

for i in range(3):
    ax[i].set_aspect('equal')
    

# Add and format colorbar 
fig.colorbar(CS, ax=ax[:], location='bottom', format="%.3f", ticks=ticks)
<matplotlib.colorbar.Colorbar at 0x7fe61bdcae20>

Sharing Axes

# Create plots
fig, ax = plt.subplots(1,3, figsize=(9,4), sharex=True, sharey=True) # Add figsize argument here to set figure size

ax[0].contourf(x, y, px_values, vmin=pmin, vmax=pmax, levels=levels, cmap="RdBu")
ax[1].contourf(x, y, py_values, vmin=pmin, vmax=pmax, levels=levels, cmap="RdBu")
CS= ax[2].contourf(x, y, pz_values, vmin=pmin, vmax=pmax, levels=levels, cmap="RdBu")

for i in range(3):
    ax[i].set_aspect('equal')
    

# Add and format colorbar 
fig.colorbar(CS, ax=ax[:], location='bottom', format="%.3f", ticks=ticks,  label=r"$\psi$")

fig.savefig("p_orbitals.png", dpi=150)