How to constrain the polar heatmap chart along the axes?

Question

I use polar heatmap im matplotlib:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from energetics import compute_s