#!/usr/bin/env python3

"""
plot of the angular part of the dz2 AO
here theta is the polar angle, not the theta angle of spheric coordinate
"""
import numpy as np
import matplotlib.pyplot as plt


def Y20(t):
    """ Compute the Y20 values """
    return np.sqrt(5. / (16. * np.pi)) * (3. * (np.cos(t))**2 - 1.)


step = 1
theta = np.radians(np.arange(0, 360, step))
Y20values = np.abs(Y20(theta))

#
# Compute nodal surface angle
#
nodalAngle = np.arccos(1. / np.sqrt(3))

#
# plot the AO and nodal surfaces
#
ax = plt.subplot(111, projection="polar")
ax.plot(theta, Y20values)

ax.plot(2 * [nodalAngle], [0, 0.7], "r", 2 * [nodalAngle + np.pi], [0, 0.7], "r",
        2 * [-nodalAngle], [0, 0.7], "r", 2 * [-nodalAngle + np.pi], [0, 0.7], "r")
ax.set_theta_zero_location("N")

plt.show()