from sympy import *
from spb import *
import numpy as np
r, theta = symbols("r, theta")
expr = cos(r**2) * exp(-r / 3)
plot3d(expr, (r, 0, 5), (theta, 1.6 * pi, 2 * pi),
    backend=KB, is_polar=True, legend=True, grid=False,
    use_cm=True, color_func=lambda x, y, z: np.sqrt(x**2 + y**2),
    wireframe=True, wf_n1=30, wf_n2=10,
    wf_rendering_kw={"width": 0.005})