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