Función Sombrero Mexicano. Hecho en Manim

Función Sombrero Mexicano

from manim import *
class funcion3d1(ThreeDScene):
    def construct(self):
        purp = "#37265A"
        rosa = "#C12786"
        rojo = "#CC0000"
        gris = "#AAAAAA"
        azul = "#0000FF"
        Azul = "#13316D"
        verde = "#008000"
        cian = "#00FFFF"
        cielo = "#C4DDFF"
        azul = "#001D6E"
        rojo = "#B20600"
        self.camera.background_color = cielo
        titulo = Text("Mexican hat", color=purp).move_to(np.array([0,0, 3.5])).scale(0.7).flip(np.array([0, 1, 1]))
        titulo2 = Tex("$z=(1-x^2-y^2)e^{-(x^2+y^2)/2}$", color=rojo).move_to(np.array([0, 0,2.8])).scale(0.7).flip(np.array([0, 1, 1]))
        self.add(titulo, titulo2)
        self.set_camera_orientation(phi=65 * DEGREES, theta=100 * DEGREES ,zoom=0.9)
        self.begin_ambient_camera_rotation(rate=0.05)
        plano = Surface(lambda u, v: np.array([v * np.cos(u),v * np.sin(u), 0]),
                        u_range=[0, 2 * PI],
                        v_range=[0.5, 3],
                        checkerboard_colors=[RED_D, RED_E],
                        resolution=(36, 32))
        plano.set_fill_by_value(
            axes=axes, colors=[(azul, -1), (GREEN, 0), (rojo, 1)], axis=2
        )

        curva2 = Surface(lambda u, v: np.array([v * np.cos(u),v * np.sin(u),1 / (v ** 2)]),
                        u_range=[0, 2 * PI],
                        v_range=[0.5, 3],
                        checkerboard_colors=[RED_D, RED_E],
                        resolution=(36, 32))
        curva2.set_fill_by_value(
            axes=axes, colors=[(azul, -1), (GREEN, 0), (rojo, 1)], axis=2
        )
        self.play(Write(plano))
        self.wait()
        self.play(Transform(plano,curva2))
        self.wait(3)